Update app.py

app.py

@@ -18,17 +18,10 @@ logger.info("Dataset loaded successfully")
 
 # Load model and tokenizer
 logger.info("Loading the model and tokenizer")
-model_name = "facebook/bart-large-cnn"
+model_name = "deepset/roberta-large-squad2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForQuestionAnswering.from_pretrained(model_name)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-logger.info(f"Model and tokenizer loaded successfully. Using device: {device}")
-
-# Preprocess the dataset
-logger.info("Preprocessing the dataset")
-context = " ".join([item.get('Text', '') for item in ds['train']])
-logger.info(f"Combined context length: {len(context)} characters")
+logger.info("Model and tokenizer loaded successfully")
 
 def clean_answer(answer):
     special_tokens = set(tokenizer.all_special_tokens)
@@ -38,56 +31,26 @@ def clean_answer(answer):
 def answer_question(question):
     logger.info(f"Received question: {question}")
     try:
-        # 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(device) for k, v in inputs.items()}
-            
-            with torch.no_grad():
-                outputs = model(**inputs)
-
-            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()))
-
-            # Break if we have a good answer
-            if score > 10:  # Adjust this threshold as needed
-                break
-
-        # 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("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=514, 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?"
         logger.info("Answer generated successfully")
-        return best_answer
+        return 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."
@@ -134,7 +97,7 @@ iface = gr.Interface(
 # Mount Gradio app to FastAPI
 app = gr.mount_gradio_app(app, iface, path="/")
 
-# For Hugging Face Spaces
+# For local development and testing
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)