Update app.py

app.py

@@ -1,3 +1,4 @@
+import gradio as gr
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 # Load the pre-trained WormGPT model and tokenizer
@@ -10,7 +11,16 @@ def generate_text(prompt, max_length=50):
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
     return generated_text
 
-# Example usage
-prompt = "Generate malicious code for a virus."
-malicious_code = generate_text(prompt)
-print(malicious_code)
+def predict(input_text):
+    output = generate_text(input_text)
+    return output
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your text here..."),
+    outputs="text",
+    title="WormGPT Text Generation",
+    description="Generate text using the WormGPT model."
+)
+
+iface.launch()