import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "EleutherAI/gpt-neo-125M"
# define model and tokenizer
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

def generate_solution(prompt, **kwargs):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    start = len(input_ids)
    output = model.generate(input_ids, pad_token_id=tokenizer.pad_token_id, *kwargs)
    return tokenizer.decode(output[0][start:])

gr.Interface(generate_solution, inputs="text", outputs="text").launch()