Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+from transformers import BartForQuestionAnswering, BartTokenizer
 import torch
 import logging
 from fastapi import FastAPI, HTTPException
@@ -18,9 +18,9 @@ logger.info("Dataset loaded successfully")
 
 # Load model and tokenizer
 logger.info("Loading the model and tokenizer")
-model_name = "deepset/roberta-large-squad2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+model_name = "facebook/bart-large-cnn"
+tokenizer = BartTokenizer.from_pretrained(model_name)
+model = BartForQuestionAnswering.from_pretrained(model_name)
 logger.info("Model and tokenizer loaded successfully")
 
 def clean_answer(answer):
@@ -34,23 +34,57 @@ def answer_question(question):
         logger.info("Combining text from dataset")
         context = " ".join([item.get('Text', '') for item in ds['train']])
         logger.info(f"Combined context length: {len(context)} characters")
-        logger.info("Tokenizing input")
-        inputs = tokenizer.encode_plus(question, context, return_tensors="pt", max_length=2048, truncation=True)
-        logger.info(f"Input tokens shape: {inputs['input_ids'].shape}")
-        logger.info("Getting model output")
-        outputs = model(**inputs)
-        logger.info(f"Output logits shapes: start={outputs.start_logits.shape}, end={outputs.end_logits.shape}")
-        logger.info("Processing output to get answer")
-        answer_start = torch.argmax(outputs.start_logits)
-        answer_end = torch.argmax(outputs.end_logits) + 1
-        raw_answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
-        answer = clean_answer(raw_answer)
-        logger.info(f"Generated answer: {answer}")
-        if not answer:
-            logger.warning("Generated answer was empty after cleaning")
-            answer = "I'm sorry, but I couldn't find a specific answer to that question based on the Bhagavad Gita. Could you please rephrase your question or ask about one of the core concepts like dharma, karma, bhakti, or the different types of yoga discussed in the Gita?"
+
+        # Implement sliding window approach
+        max_length = 1024
+        stride = 512
+        answers = []
+        for i in range(0, len(context), stride):
+            chunk = context[i:i+max_length]
+            
+            inputs = tokenizer.encode_plus(
+                question, 
+                chunk, 
+                return_tensors="pt", 
+                max_length=max_length,
+                truncation=True,
+                padding='max_length'
+            )
+            
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+            
+            logger.info(f"Input tokens shape: {inputs['input_ids'].shape}")
+            
+            logger.info("Getting model output")
+            with torch.no_grad():
+                outputs = model(**inputs)
+
+            logger.info(f"Output logits shapes: start={outputs.start_logits.shape}, end={outputs.end_logits.shape}")
+            
+            answer_start = torch.argmax(outputs.start_logits)
+            answer_end = torch.argmax(outputs.end_logits) + 1
+
+            ans = tokenizer.convert_tokens_to_string(
+                tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end])
+            )
+            
+            score = torch.max(outputs.start_logits) + torch.max(outputs.end_logits)
+            answers.append((ans, score.item()))
+
+        # Select best answer
+        best_answer = max(answers, key=lambda x: x[1])[0]
+        
+        # Post-processing
+        best_answer = clean_answer(best_answer)
+        best_answer = best_answer.capitalize()
+        
+        logger.info(f"Generated answer: {best_answer}")
+        if not best_answer or len(best_answer) < 5:
+            logger.warning("Generated answer was empty or too short after cleaning")
+            best_answer = "I'm sorry, but I couldn't find a specific answer to that question based on the Bhagavad Gita. Could you please rephrase your question or ask about one of the core concepts like dharma, karma, bhakti, or the different types of yoga discussed in the Gita?"
+        
         logger.info("Answer generated successfully")
-        return answer
+        return best_answer
     except Exception as e:
         logger.error(f"Error in answer_question function: {str(e)}")
         return "I'm sorry, but an error occurred while processing your question. Please try again later."