import gradio as gr
from transformers import RobertaTokenizer, T5ForConditionalGeneration

tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-small')
model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-small')

def generate_code(text):
    input_ids = tokenizer(text, return_tensors="pt").input_ids
    generated_ids = model.generate(input_ids, max_length=10)
    return tokenizer.decode(generated_ids[0], skip_special_tokens=True)

examples = [
    ["def greet(user): print(f'hello <extra_id_0>!')"],
]

iface = gr.Interface(fn=generate_code, inputs="text", outputs="text", examples=examples)
iface.launch()