Update app.py

app.py

@@ -1,26 +1,31 @@
 import gradio as gr
 from transformers import AutoModelForQuestionAnswering, AutoTokenizer
 import torch
+import torch.nn.functional as F
 
 # Load model and tokenizer
-MODEL_NAME = "your-hf-username/raft-qa"
+MODEL_NAME = "S-Dreamer/raft-qa-space"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForQuestionAnswering.from_pretrained(MODEL_NAME)
 
 def answer_question(context, question):
-    inputs = tokenizer(question, context, return_tensors="pt", truncation=True, max_length=512)
+    inputs = tokenizer(
+        question, context, return_tensors="pt", truncation=True, max_length=512, stride=128, return_overflowing_tokens=True
+    )
     with torch.no_grad():
         outputs = model(**inputs)
     
-    start_scores, end_scores = outputs.start_logits, outputs.end_logits
-    start_idx = torch.argmax(start_scores)
-    end_idx = torch.argmax(end_scores) + 1
-    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start_idx:end_idx]))
+    start_probs = F.softmax(outputs.start_logits, dim=-1)
+    end_probs = F.softmax(outputs.end_logits, dim=-1)
+    start_idx = torch.argmax(start_probs)
+    end_idx = torch.argmax(end_probs) + 1
+
+    answer = tokenizer.decode(inputs["input_ids"][0][start_idx:end_idx], skip_special_tokens=True)
     
     return answer if answer.strip() else "No answer found."
 
 # Define UI
-with gr.Blocks(theme="soft") as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# 🤖 RAFT: Retrieval-Augmented Fine-Tuning for QA")
     gr.Markdown("Ask a question based on the provided context and see how RAFT improves response accuracy!")