import streamlit as st
import torch  # <-- Add this import
from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline

# Load the model and tokenizer in full precision and ensure it's on CPU
model = AutoModelForCausalLM.from_pretrained("PAIXAI/Astrid-1B")
tokenizer = AutoTokenizer.from_pretrained("PAIXAI/Astrid-1B")

# Initialize the pipeline with the model and tokenizer to run on CPU
generate_text = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=-1)  # -1 forces CPU usage

# Streamlit UI
st.title("Astrid-1B Chatbot")
st.write("Test the Astrid-1B chatbot from Hugging Face!")

user_input = st.text_input("Enter your question:")
if user_input:
    try:
        response = generate_text(user_input, min_new_tokens=2, max_new_tokens=256, do_sample=False, num_beams=1, temperature=0.3, repetition_penalty=1.2, renormalize_logits=True)
        st.write("Response:", response[0]["generated_text"])
    except Exception as e:
        st.write("Error:", str(e))

st.write("Note: This is a simple UI for demonstration purposes. Ensure you have the required libraries installed and adjust the model parameters as needed.")