Update app.py

app.py

@@ -9,14 +9,15 @@ from retriever import retrieve_documents
 #os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 
 # Load Mistral 7B model
-MODEL_NAME = "mistralai/Mistral-7B-v0.1"
+#MODEL_NAME = "mistralai/Mistral-7B-v0.1"
+MODEL_NAME = "microsoft/phi3-mini-4k-instruct"
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"), cache_dir="/tmp/huggingface")
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, 
                                              use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"), 
                                              cache_dir="/tmp/huggingface", 
                                              device_map="auto",
-                                             torch_dtype=torch.float16,
-                                             load_in_4bit=True
+                                             torch_dtype=torch.float16
                                             ) 
 
 # Create inference pipeline
@@ -27,7 +28,7 @@ app = FastAPI()
 
 @app.get("/")
 def read_root():
-    return {"message": "Mistral 7B RAG API is running!"}
+    return {"message": "Phi3 Mini RAG API is running!"}
 
 @app.get("/generate/")
 def generate_response(query: str = Query(..., title="User Query")):