import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

model_path = "./fine_tuned_gpt_neo"  # Шлях до донавченої моделі
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

def generate_response(prompt, max_length=100, temperature=0.5, top_p=0.9, top_k=50):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
      inputs["input_ids"],
      max_length=max_length,
      temperature=temperature,
      top_p=top_p,
      top_k=top_k,
      do_sample=True,
      pad_token_id=tokenizer.eos_token_id
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

interface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=2, placeholder="Введіть запит про космос..."),
        gr.Slider(50, 500, step=10, label="Максимальна довжина тексту"),
        gr.Slider(0.1, 1.0, step=0.1, label="Температура (креативність)"),
        gr.Slider(0.1, 1.0, step=0.1, label="Top-p (фільтрація найімовірніших слів)"),
    ],
    outputs=gr.Textbox(label="Згенерований текст"),
    title="GPT-Neo Чат-бот про космос",
    description="Цей чат-бот використовує GPT-Neo, донавчений на текстах про космос. Введіть запит, і бот згенерує текст.",
)

interface.launch()