import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("LoewolfAI/L-GPT_1.5")
model = AutoModelForCausalLM.from_pretrained("LoewolfAI/L-GPT_1.5")

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

    max_new_tokens = 90
    output = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_new_tokens=max_new_tokens,
        num_beams=5,
        no_repeat_ngram_size=2,
        early_stopping=True,
        temperature=0.7,
        top_p=0.9,
        top_k=50,
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )

    text = tokenizer.decode(output[0], skip_special_tokens=True)
    return text


css = """
h1 { text-align: center; }
#duplicate-button {
  margin: auto;
  color: white;
  background: #1565c0;
  border-radius: 100vh;
}
.contain {
  max-width: 900px;
  margin: auto;
  padding-top: 1.5rem;
}
"""

iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=2, placeholder="Nachricht eingeben...", label="Deine Nachricht"),
    outputs=gr.Textbox(label="Löwolf Chat Antwort", placeholder="Antwort erscheint hier...", interactive=False, lines=10),
    title="Löwolf Chat",
    css=css,
)

iface.launch()