Update app.py

app.py

@@ -10,15 +10,15 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 trained_model.to(device)
 untrained_model.to(device)
 
-def generate(commentary_text):
-    # Generate text using the trained model
+def generate(commentary_text, max_length, temperature):
+    # Generate text using the finetuned model
     input_ids = trained_tokenizer(commentary_text, return_tensors="pt").input_ids.to(device)
-    trained_output = trained_model.generate(input_ids, max_length=60, num_beams=5, do_sample=False)
+    trained_output = trained_model.generate(input_ids, max_length=max_length, num_beams=5, do_sample=False, temperature=temperature)
     trained_text = trained_tokenizer.decode(trained_output[0], skip_special_tokens=True)
     
-    # Generate text using the untrained model
+    # Generate text using the base model
     input_ids = untrained_tokenizer(commentary_text, return_tensors="pt").input_ids.to(device)
-    untrained_output = untrained_model.generate(input_ids, max_length=60, num_beams=5, do_sample=False)
+    untrained_output = untrained_model.generate(input_ids, max_length=max_length, num_beams=5, do_sample=False,temperature=temperature)
     untrained_text = untrained_tokenizer.decode(untrained_output[0], skip_special_tokens=True)
     
     return trained_text, untrained_text
@@ -26,10 +26,17 @@ def generate(commentary_text):
 # Create Gradio interface
 iface = gr.Interface(
     fn=generate, 
-    inputs="text", 
-    outputs=["text", "text"],
+    inputs=[
+        gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),
+        gr.inputs.Slider(minimum=10, maximum=100, default=50, label="Max Length"),        
+        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.7, label="Temperature")
+    ], 
+    outputs=[
+        gr.outputs.Textbox(label="commentary generation from finetuned GPT2 Model"), 
+        gr.outputs.Textbox(label="commentary generation from base GPT2 Model")
+    ],
     title="GPT-2 Text Generation",
-    description="start writing a cricket commentary and GPT-2 will continue it using both a trained and untrained model."
+    description="start writing a cricket commentary and GPT-2 will continue it using both a finetuned and base model."
 )
 
 # Launch the app