Update app.py

app.py

@@ -10,10 +10,6 @@ def load_model(model_name):
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 
-# Models to compare
-original_model_name = "Vishwas1/hummingbird-base-marathi-finetuned"  
-fine_tuned_model_name = "Vishwas1/hummingbird-base-marathi-finetuned-finetuned"  
-
 # Load Hugging Face token
 hf_token = os.getenv('HF_API_TOKEN')
 if not hf_token:
@@ -22,16 +18,16 @@ if not hf_token:
 # Login to Hugging Face Hub
 login(hf_token)
 
-# 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)
+# Function to compare text generation from both models
+def compare_models(prompt, original_model_name, fine_tuned_model_name):
+    # Load the original and fine-tuned models based on user input
+    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()
+    # Ensure models are in evaluation mode
+    original_model.eval()
+    fine_tuned_model.eval()
 
-# 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():
@@ -54,13 +50,15 @@ def compare_models(prompt):
 # Gradio Interface
 iface = gr.Interface(
     fn=compare_models,
-    inputs=gr.Textbox(lines=5, placeholder="Enter text here...", label="Input Text"),
+    inputs=[
+        gr.Textbox(lines=5, placeholder="Enter text here...", label="Input Text"),
+        gr.Textbox(lines=1, placeholder="Enter original model name...", label="Original Model Name", default="Vishwas1/hummingbird-base-marathi-finetuned"),
+        gr.Textbox(lines=1, placeholder="Enter fine-tuned model name...", label="Fine-Tuned Model Name", default="Vishwas1/hummingbird-base-marathi-finetuned-finetuned")
+    ],
     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."
+    description="Enter a prompt and model names to generate text from the original and fine-tuned models."
 )
 
 iface.launch()
 
-
-