Update app.py

app.py

@@ -96,6 +96,9 @@ def generate_answer(question, context, ground, do_pretrained):
         max_length=max_length,
     ).input_ids
 
+    # Decode the context back
+    decoded_context = tokenizer.decode(input_ids, skip_special_tokens=True)[len(f"question: {question} context: "):]
+
     # Generate the answer
     with torch.no_grad():
         generated_ids = model.generate(input_ids=input_ids, max_new_tokens=max_target_length)
@@ -110,7 +113,7 @@ def generate_answer(question, context, ground, do_pretrained):
             pretrained_generated_ids = pretrained_model.generate(input_ids=input_ids, max_new_tokens=max_target_length)
             pretrained_answer = tokenizer.decode(pretrained_generated_ids[0], skip_special_tokens=True)
 
-    return generated_answer, context, pretrained_answer
+    return generated_answer, decoded_context, pretrained_answer
 
 
 # Define a function to list examples from the dataset