Update app.py

app.py

@@ -1,37 +1,36 @@
-import torch
 import gradio as gr
-from transformers import AutoModelForQuestionAnswering, AutoTokenizer
+import torch
+from fastai.text.all import load_learner
+from huggingface_hub import hf_hub_download
 
-# Load the model and tokenizer from Hugging Facehh
-model = AutoModelForQuestionAnswering.from_pretrained("rahul7star/fastai-rahul-text-model-v02")
-tokenizer = AutoTokenizer.from_pretrained("rahul7star/fastai-rahul-text-model-v02")
+# Step 8: Download the model from Hugging Face and load it
+def load_model():
+    try:
+        # Download the model .pth file from Hugging Face
+        model_path = hf_hub_download(
+            repo_id="rahul7star/fastai-rahul-text-model-v02", 
+            filename="model.pth"
+        )
+        
+        # Load the model using FastAI's load_learner method
+        learn = load_learner(model_path)
+        print("Model loaded successfully from Hugging Face.")
+        return learn
+    except Exception as e:
+        print(f"Error loading the model: {e}")
+        return None
 
-# 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
+# Load the model
+learn = load_model()
 
-# 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
-)
+# Step 9: Define the Gradio Interface
+def predict(input_text):
+    try:
+        # Get prediction from the model
+        pred, _, probs = learn.predict(input_text)
+        return f"Prediction: {pred}, Confidence: {probs.max():.2f}"
+    except Exception as e:
+        return f"Error during prediction: {e}"
 
-# Launch the interface
-interface.launch()
+# Step 10: Create Gradio Interface
+gr.Interface(fn=predict, inputs="text", outputs="text").launch()