import gradio as gr
from transformers import TFGPT2LMHeadModel, GPT2Tokenizer
description = "Chatbots for customer service with GPT-2"
title = "Please enter the problem you are experiencing"
examples = [["I didn't get my order delivered for a long time."]]

from transformers import TFGPT2LMHeadModel, GPT2Tokenizer

model_name_or_path = "Jinpkk/codeparrot-ds"  

tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path)
model = TFGPT2LMHeadModel.from_pretrained(model_name_or_path)


def generate_response(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors='tf')
    beam_output = model.generate(
        input_ids,
        max_length=128,
        num_beams=5,
        no_repeat_ngram_size=2,
        early_stopping=True,
        pad_token_id=tokenizer.eos_token_id
    )
    generated_text = tokenizer.decode(beam_output[0], skip_special_tokens=True)
    generated_text = generated_text.split('[PAD]')[0].strip()
    generated_text = generated_text.split('^')[0].strip()
    generated_text = generated_text.replace(input_text.strip(), '').strip()
    return generated_text

    

interface = gr.Interface(
            fn=generate_response,
            inputs=gr.inputs.Textbox(lines=5, placeholder="Enter your text..."),
            outputs=gr.outputs.Textbox(),
            description=description,
            title=title,
            examples=examples
)


interface.launch()