Update app.py

app.py

@@ -30,7 +30,10 @@ def preprocess_text(text):
 
 # ====================== MODELS ======================
 classifier_model = "Ginidu2003/Distilbert-Base-News-classifier"
-qa_model_name = "deepset/roberta-base-squad2"
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+
+qa_tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
+qa_model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
 
 # Load QA pipeline using a supported method
 qa_pipeline = pipeline("document-question-answering", model=qa_model_name, device=-1)
@@ -75,13 +78,27 @@ def classify_csv(file):
 def answer_question(news_content, question):
     if not news_content.strip() or not question.strip():
         return "Please enter both news content and a question."
+    
     try:
-        result = qa_pipeline(question=question, context=news_content)
-        answer = result[0]['answer'] if isinstance(result, list) else result['answer']
-        score = result[0]['score'] if isinstance(result, list) else result['score']
-        return f"**Answer:** {answer}\n\n**Confidence:** {score:.2%}"
+        inputs = qa_tokenizer(question, news_content, return_tensors="pt", truncation=True, max_length=512)
+        
+        with torch.no_grad():
+            outputs = qa_model(**inputs)
+        
+        start_scores = outputs.start_logits
+        end_scores = outputs.end_logits
+        
+        start_idx = torch.argmax(start_scores)
+        end_idx = torch.argmax(end_scores) + 1
+        
+        answer = qa_tokenizer.decode(inputs.input_ids[0][start_idx:end_idx])
+        
+        confidence = torch.max(torch.softmax(start_scores, dim=1)).item()
+        
+        return f"**Answer:** {answer}\n\n**Confidence:** {confidence:.2%}"
+    
     except Exception as e:
-        return f"Error processing question: {str(e)}"
+        return f"Error: {str(e)}"
 
 # ====================== GRADIO INTERFACE ======================
 with gr.Blocks(title="Daily Mirror News Classifier") as demo: