Update app.py

app.py

@@ -1,16 +1,17 @@
 import gradio as gr
 from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
 import json
+import torch
 
 # Define hyperparameters
-learning_rate = 3e-5
-batch_size = 16
-epochs = 3
-max_seq_length = 512
-warmup_steps = 100
-weight_decay = 0.01
-dropout_prob = 0.1
-gradient_clip_value = 1.0
+learning_rate = 3e-5  # Slightly lower learning rate
+batch_size = 8  # Smaller batch size to allow for more precise updates
+epochs = 4  # Slightly more training epochs
+max_seq_length = 256  # Smaller sequence length, especially if the majority of your questions and contexts are shorter
+warmup_steps = 200  # Longer warmup phase
+weight_decay = 0.01  # Keep weight decay as it is
+dropout_prob = 0.2  # Slightly higher dropout for regularization
+gradient_clip_value = 1.0  # Keep gradient clip value as it is
 
 context_val = ''
 
@@ -35,10 +36,19 @@ def q_n_a_fn(context, text):
     with torch.no_grad():
         outputs = q_n_a_model(**inputs)
     
-    # Decode and return the answer
+    # Get the predicted answer span indices
     start_idx, end_idx = torch.argmax(outputs.start_logits), torch.argmax(outputs.end_logits)
-    answer = tokenizer.decode(inputs["input_ids"][0][start_idx:end_idx+1])
-    
+
+    # Ensure indices are within bounds
+    start_idx = min(start_idx, len(inputs["input_ids"][0]) - 1)
+    end_idx = min(end_idx, len(inputs["input_ids"][0]) - 1)
+
+    # Find the answer tokens in the input
+    answer_tokens = inputs["input_ids"][0][start_idx : end_idx + 1]
+
+    # Decode the answer tokens into a human-readable answer
+    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
+
     return answer
 
 def classification_fn(text):
@@ -75,4 +85,4 @@ with gr.Blocks(theme='gradio/soft') as demo:
             gr.Interface(fn=classification_fn, inputs=[context], outputs="text")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()