Update app.py

app.py

@@ -1,18 +1,12 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline
 
-
-# Load the model in full precision
+# Load the model and tokenizer in full precision
 model = AutoModelForCausalLM.from_pretrained("PAIXAI/Astrid-1B").to(dtype=torch.float32)
+tokenizer = AutoTokenizer.from_pretrained("PAIXAI/Astrid-1B")
 
-
-# Initialize the model and tokenizer for CPU
-generate_text = pipeline(
-    model="PAIXAI/Astrid-1B",
-    torch_dtype="auto",
-    trust_remote_code=True,
-    use_fast=True
-)
+# Initialize the pipeline with the model and tokenizer
+generate_text = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
 # Streamlit UI
 st.title("Astrid-1B Chatbot")