app.py

import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load the pre-trained WormGPT model and tokenizer
model = GPT2LMHeadModel.from_pretrained("wormgpt")
tokenizer = GPT2Tokenizer.from_pretrained("wormgpt")

def generate_text(prompt, max_length=50):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text

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()