Update app.py

app.py

@@ -1,20 +1,37 @@
+import torch
 import gradio as gr
-from transformers import BertForQuestionAnswering, BertTokenizer, pipeline
-model_name = "rahul7star/Rahul-Context-AI-01"
-# Load your fine-tuned model and tokenizer
-model = BertForQuestionAnswering.from_pretrained(model_name)  # Path to the fine-tuned model
-tokenizer = BertTokenizer.from_pretrained(model_name)  # Path to the tokenizer
+from transformers import AutoModelForQuestionAnswering, AutoTokenizer
 
-# Initialize the custom pipeline with your fine-tuned model
-qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
+# Load the model and tokenizer from Hugging Face
+model = AutoModelForQuestionAnswering.from_pretrained("rahul7star/fastai-rahul-text-model-v02")
+tokenizer = AutoTokenizer.from_pretrained("rahul7star/fastai-rahul-text-model-v02")
 
-# Function to get the answer using the fine-tuned model
-def answer_question(context, question):
-    result = qa_pipeline({'context': context, 'question': question})
-    return result['answer']
+# Function to handle predictions (for question-answering tasks)
+def get_answer(question):
+    # Tokenize the input question
+    inputs = tokenizer(question, return_tensors="pt")
+    
+    # Get model prediction (start and end positions for the answer)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Extract start and end positions of the predicted answer
+    start_idx = torch.argmax(outputs.start_logits)
+    end_idx = torch.argmax(outputs.end_logits)
+    
+    # Convert the token IDs back to text
+    answer_tokens = inputs.input_ids[0][start_idx:end_idx+1]
+    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
+    
+    return answer
 
-# Define the Gradio interface
-interface = gr.Interface(fn=answer_question, inputs=["text", "text"], outputs="text")
+# Set up the Gradio interface
+interface = gr.Interface(
+    fn=get_answer,  # Function to call for inference
+    inputs=gr.Textbox(label="Ask a Question"),  # Input field for question
+    outputs=gr.Textbox(label="Answer"),  # Output field for the model's answer
+    live=True  # Set to True for real-time interaction
+)
 
 # Launch the interface
 interface.launch()