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

# Initialisierung des Modells und des Tokenizers
tokenizer = GPT2Tokenizer.from_pretrained("Loewolf/GPT_1")
model = GPT2LMHeadModel.from_pretrained("Loewolf/GPT_1")

def generate_text(prompt):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    attention_mask = torch.ones(input_ids.shape, dtype=torch.bool)

    max_length = model.config.n_positions if len(input_ids[0]) > model.config.n_positions else len(input_ids[0]) + 50
    beam_output = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_length=max_length,
        min_length=4,
        num_beams=5,
        no_repeat_ngram_size=2,
        early_stopping=True,
        temperature=0.9,
        top_p=0.90,
        top_k=50,
        length_penalty=2.0,
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id 
    )
    
    text = tokenizer.decode(beam_output[0], skip_special_tokens=True)
    return text

DESCRIPTION = """\
#Löwolf GPT1 Chat
"""
css = """
    body { font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif; }
    .gradio_app { display: flex; flex-direction: column-reverse; max-width: 800px; margin: auto; }
    .gradio_interface { box-shadow: 0 0 20px rgba(0,0,0,0.1); }
    .gradio_text_input { border-radius: 20px; }
    .gradio_text_output { height: calc(100vh - 150px); border-radius: 20px; overflow-y: auto; }
    button { border-radius: 20px; }
    #input-box { order: 2; }
    #output-box { order: 1; }
    #input { height: 50px; }
    #output { height: calc(100vh - 150px); }
"""

# Update the Blocks structure to reflect the desired layout
with gr.Blocks(css=css) as demo:
    gr.Markdown(DESCRIPTION)
    
    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
    with gr.Row().style(equal_height=False):
        output_box = gr.Textbox(interactive=False, lines=25, placeholder="Responses will appear here...", elem_id="output-box")
        input_box = gr.Textbox(lines=1, placeholder="Type a message...", elem_id="input-box")
    generate_button = gr.Button("Submit")
    generate_button.click(
        fn=generate,
        inputs=[input_box],
        outputs=output_box
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()