Update app.py

app.py

@@ -51,8 +51,18 @@ print("✅ ChromaDB initialized!")
 
 
 # 🔹 Function to Generate LLM Responses
+import torch
+
+# 🔹 Detect Device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"✅ Using device: {device}")
+
 def generate_humanized_response(query, retrieved_text):
     """Passes retrieved chunks through Mistral-7B to improve readability."""
+    
+    # 🔹 Truncate retrieved text to avoid long input errors
+    retrieved_text = retrieved_text[:500]  
+
     prompt = f"""You are a medical assistant. Answer the following question based on retrieved text:
 
     Retrieved Text:
@@ -63,13 +73,12 @@ def generate_humanized_response(query, retrieved_text):
     Provide a well-structured, human-like response:
     """
     
-    inputs = llm_tokenizer(prompt, return_tensors="pt").to("cuda")
-    output = llm_model.generate(**inputs, max_new_tokens=150, do_sample=True, stream=True)  
+    inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)  # ✅ Uses GPU if available, otherwise CPU
+    output = llm_model.generate(**inputs, max_new_tokens=150, do_sample=True)  
     response = llm_tokenizer.decode(output[0], skip_special_tokens=True)
     
     return response
-from transformers import AutoModel, AutoTokenizer
-import torch
+
 
 # 🔹 Load BioMedBERT for Embeddings
 embed_model_name = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract"