Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from huggingface_hub import login
 import os
@@ -16,20 +16,21 @@ fine_tuned_model_name = "Vishwas1/hummingbird-finetuned-marathi"  # Replace with
 
 # Load models
 hf_token = os.getenv('HUGGINGFACE_TOKEN')
-    if not hf_token:
-        return "Error: Hugging Face token not found. Please set it as a secret."
-    
+if not hf_token:
+    raise ValueError("Error: Hugging Face token not found. Please set it as an environment variable.")
+
 # Login to Hugging Face Hub
 login(hf_token)
-original_tokenizer = AutoTokenizer.from_pretrained(original_model_name)
-original_model = AutoModelForSequenceClassification.from_pretrained(original_model_name)
-fine_tuned_tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model_name)
-fine_tuned_model = AutoModelForSequenceClassification.from_pretrained(fine_tuned_model_name)
+
+# Load the original and fine-tuned models
+original_tokenizer, original_model = load_model(original_model_name)
+fine_tuned_tokenizer, fine_tuned_model = load_model(fine_tuned_model_name)
 
 # Ensure models are in evaluation mode
 original_model.eval()
 fine_tuned_model.eval()
 
+# Function to compare models
 def compare_models(text):
     # Original model prediction
     inputs_orig = original_tokenizer(text, return_tensors='pt', truncation=True, padding=True)
@@ -68,3 +69,5 @@ iface = gr.Interface(
 )
 
 iface.launch()
+
+