Update app.py

app.py

@@ -3,16 +3,16 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
 import faiss
 
-# Use a lightweight CPU-friendly model
-model_id = "tiiuae/falcon-rw-1b"
+# Use a smaller model for testing
+model_id = "mistralai/Mistral-7B-Instruct-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
 
-# Embedding model (small but good)
+# Embedding model
 embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-# In-memory document store
+# Storage
 doc_texts = []
 index = None
 
@@ -62,3 +62,4 @@ with gr.Blocks() as demo:
     chatbot.submit(ask_bot, inputs=chatbot, outputs=output)
 
 demo.launch()
+