Update app.py

app.py

@@ -4,7 +4,7 @@ import torch
 from huggingface_hub import login
 import os
 
-# Load the original pre-trained model
+# Load text generation model
 def load_model(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
@@ -12,9 +12,9 @@ def load_model(model_name):
 
 # Models to compare
 original_model_name = "Vishwas1/hummingbird-base-marathi"  # Replace with your original model
-fine_tuned_model_name = "Vishwas1/hummingbird-finetuned-marathi"  # Replace with your fine-tuned model's repo ID
+fine_tuned_model_name = "Vishwas1/hummingbird-finetuned-marathi"  # Replace with your fine-tuned model
 
-# Load models
+# Load Hugging Face token
 hf_token = os.getenv('HF_API_TOKEN')
 if not hf_token:
     raise ValueError("Error: Hugging Face token not found. Please set it as an environment variable.")
@@ -30,32 +30,24 @@ fine_tuned_tokenizer, fine_tuned_model = load_model(fine_tuned_model_name)
 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)
+# Function to compare text generation from both models
+def compare_models(prompt):
+    # Generate text with the original model
+    inputs_orig = original_tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
-        outputs_orig = original_model(**inputs_orig)
-    logits_orig = outputs_orig.logits
-    probs_orig = torch.softmax(logits_orig, dim=1)
-    pred_orig = torch.argmax(probs_orig, dim=1).item()
-    confidence_orig = probs_orig[0][pred_orig].item()
+        generated_ids_orig = original_model.generate(inputs_orig["input_ids"], max_length=100)
+    generated_text_orig = original_tokenizer.decode(generated_ids_orig[0], skip_special_tokens=True)
 
-    # Fine-tuned model prediction
-    inputs_fine = fine_tuned_tokenizer(text, return_tensors='pt', truncation=True, padding=True)
+    # Generate text with the fine-tuned model
+    inputs_fine = fine_tuned_tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
-        outputs_fine = fine_tuned_model(**inputs_fine)
-    logits_fine = outputs_fine.logits
-    probs_fine = torch.softmax(logits_fine, dim=1)
-    pred_fine = torch.argmax(probs_fine, dim=1).item()
-    confidence_fine = probs_fine[0][pred_fine].item()
-
-    # Map predictions to labels (adjust based on your model's labels)
-    labels = {0: "Negative", 1: "Positive"}
+        generated_ids_fine = fine_tuned_model.generate(inputs_fine["input_ids"], max_length=100)
+    generated_text_fine = fine_tuned_tokenizer.decode(generated_ids_fine[0], skip_special_tokens=True)
 
+    # Return the generated text from both models for comparison
     result = {
-        "Original Model Prediction": f"{labels[pred_orig]} ({confidence_orig:.2f})",
-        "Fine-Tuned Model Prediction": f"{labels[pred_fine]} ({confidence_fine:.2f})"
+        "Original Model Output": generated_text_orig,
+        "Fine-Tuned Model Output": generated_text_fine
     }
     return result
 
@@ -63,11 +55,12 @@ def compare_models(text):
 iface = gr.Interface(
     fn=compare_models,
     inputs=gr.Textbox(lines=5, placeholder="Enter text here...", label="Input Text"),
-    outputs=gr.JSON(label="Model Predictions"),
-    title="Compare Original and Fine-Tuned Models",
-    description="Enter text to see predictions from the original and fine-tuned models."
+    outputs=gr.JSON(label="Generated Texts"),
+    title="Compare Text Generation from Original and Fine-Tuned Models",
+    description="Enter a prompt to generate text from the original and fine-tuned models."
 )
 
 iface.launch()
 
 
+