Update app.py

app.py

@@ -12,21 +12,30 @@ model_name_or_path = "Jinpkk/codeparrot-ds"
 tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path)
 model = TFGPT2LMHeadModel.from_pretrained(model_name_or_path)
 
+
 def generate_response(input_text):
     input_ids = tokenizer.encode(input_text, return_tensors='tf')
-    beam_output = model.generate(input_ids, max_length=150, num_beams=5, no_repeat_ngram_size=2, early_stopping=True)
-    output = tokenizer.decode(beam_output[0], skip_special_tokens=True)
-    
-    # Format the output as needed
-    formatted_output = output.replace(input_text, "").strip()  # Removing input text from output
-    generated_text = formatted_output.split('[PAD]')[0].strip()
-    return formatted_output
+    beam_output = model.generate(
+        input_ids,
+        max_length=128,
+        num_beams=5,
+        no_repeat_ngram_size=2,
+        early_stopping=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    generated_text = tokenizer.decode(beam_output[0], skip_special_tokens=True)
+    generated_text = generated_text.split('[PAD]')[0].strip()
+    generated_text = generated_text.replace(input_text.strip(), '').strip()
+    return generated_text
+
     
 
 interface = gr.Interface(
             fn=generate_response,
+            inputs=gr.inputs.Textbox(lines=5, placeholder="Enter your text..."),
+            outputs=gr.outputs.Textbox(),
             description=description,
-            title = title,
+            title=title,
             examples=examples
 )