import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

def generate_text(input_text, temperature):
    input_text = f"Prompt: {input_text}\nResponse:"
    encoded_input = tokenizer.encode(input_text, max_length=1024)
    input_ids = torch.tensor(encoded_input).unsqueeze(0)

    output = model.generate(input_ids, temperature=temperature, max_length=1024, top_k=50, top_p=1.0, repetition_penalty=1.0)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text #.strip().split("\n")[-1]

inputs = [gr.inputs.Textbox(lines=1, label="Input"),
          gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.1, default=0.7, label="Temperature")]
outputs = gr.outputs.Textbox()
interface = gr.Interface(generate_text, inputs, outputs, title="GPT-2 Text Generator")
interface.launch()