Upload app.py

app.py

@@ -158,8 +158,8 @@ emotion_classifier = hf_pipeline("text-classification", model="nateraw/bert-base
 
 device = 0 if torch.cuda.is_available() else -1
 tokenizer = AutoTokenizer.from_pretrained(READER_MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(READER_MODEL_NAME)
-#model = AutoModelForCausalLM.from_pretrained(READER_MODEL_NAME).to(device)
+#model = AutoModelForCausalLM.from_pretrained(READER_MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(READER_MODEL_NAME).to(device)
 READER_LLM = pipeline(
     model=model,
     tokenizer=tokenizer,
@@ -169,7 +169,7 @@ READER_LLM = pipeline(
     repetition_penalty=1.1,
     return_full_text=False,
     max_new_tokens=500,
-#    device=device,
+    device=device,
 )
 # -------------------------------
 # 🔧 Whisper Model Setup
@@ -245,7 +245,14 @@ def process_query(user_query, input_type="text"):
     )
 
     # Generate response
-    answer = READER_LLM(RAG_PROMPT_TEMPLATE)[0]["generated_text"]
+    #answer = READER_LLM(RAG_PROMPT_TEMPLATE)[0]["generated_text"]
+    try:
+        response = READER_LLM(RAG_PROMPT_TEMPLATE)
+        answer = response[0]["generated_text"] if response and "generated_text" in response[0] else "⚠️ No output generated."
+    except Exception as e:
+        print("❌ Error during generation:", e)
+        answer = "⚠️ An error occurred while generating the response."
+
     # Estimate severity score from token probabilities
     severity_score = round(np.random.uniform(0.6, 1.0), 2)
     answer += f"\n\n🧭 Confidence Score: {value}"