Update app.py

app.py

@@ -25,6 +25,11 @@ 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")
+
 def clean_answer(answer):
     special_tokens = set(tokenizer.all_special_tokens)
     cleaned_answer = ' '.join(token for token in answer.split() if token not in special_tokens)
@@ -33,14 +38,11 @@ def clean_answer(answer):
 def answer_question(question):
     logger.info(f"Received question: {question}")
     try:
-        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")
-
         # Implement sliding window approach
         max_length = 1024
         stride = 512
         answers = []
+        
         for i in range(0, len(context), stride):
             chunk = context[i:i+max_length]
             
@@ -55,14 +57,9 @@ def answer_question(question):
             
             inputs = {k: v.to(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
 
@@ -73,6 +70,10 @@ def answer_question(question):
             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]