Update app.py

app.py

@@ -15,6 +15,13 @@ def load_model(model_name):
         tokenizer = GPT2Tokenizer.from_pretrained(model_name)
     
     model = AutoModelForCausalLM.from_pretrained(model_name)
+    
+    # Assign eos_token as pad_token if not already set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    if model.config.pad_token_id is None:
+        model.config.pad_token_id = tokenizer.pad_token_id
+    
     return tokenizer, model
 
 # Load Hugging Face token
@@ -36,16 +43,34 @@ def compare_models(prompt, original_model_name, fine_tuned_model_name):
     fine_tuned_model.eval()
 
     # Generate text with the original model
-    inputs_orig = original_tokenizer(prompt, return_tensors="pt")
+    inputs_orig = original_tokenizer(prompt, return_tensors="pt", padding=True)
     with torch.no_grad():
-        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)
+        generated_ids_orig = original_model.generate(
+            input_ids=inputs_orig["input_ids"],
+            attention_mask=inputs_orig["attention_mask"],
+            max_length=100,
+            pad_token_id=original_tokenizer.pad_token_id
+        )
+    generated_text_orig = original_tokenizer.decode(
+        generated_ids_orig[0],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True  # Optional
+    )
 
     # Generate text with the fine-tuned model
-    inputs_fine = fine_tuned_tokenizer(prompt, return_tensors="pt")
+    inputs_fine = fine_tuned_tokenizer(prompt, return_tensors="pt", padding=True)
     with torch.no_grad():
-        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)
+        generated_ids_fine = fine_tuned_model.generate(
+            input_ids=inputs_fine["input_ids"],
+            attention_mask=inputs_fine["attention_mask"],
+            max_length=100,
+            pad_token_id=fine_tuned_tokenizer.pad_token_id
+        )
+    generated_text_fine = fine_tuned_tokenizer.decode(
+        generated_ids_fine[0],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True  # Optional
+    )
 
     # Return the generated text from both models for comparison
     result = {
@@ -59,8 +84,8 @@ iface = gr.Interface(
     fn=compare_models,
     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"),
-        gr.Textbox(lines=1, placeholder="Enter fine-tuned model name...", label="Fine-Tuned Model Name")
+        gr.Textbox(lines=1, placeholder="e.g., gpt2-medium", label="Original Model Name"),
+        gr.Textbox(lines=1, placeholder="e.g., your-username/gpt2-medium-finetuned", label="Fine-Tuned Model Name")
     ],
     outputs=gr.JSON(label="Generated Texts"),
     title="Compare Text Generation from Original and Fine-Tuned Models",