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

model_name = "Azurro/APT3-1B-Instruct-v1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
pipeline = transformers.pipeline("text-generation", model=model, tokenizer=tokenizer)


def generate_text(input_text, max_tokens, temperature, top_p):
    
    sequences = pipeline(max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, eos_token_id=tokenizer.eos_token_id, text_inputs=input_text)
    txt = ""

    for seq in sequences:
        txt += seq['generated_text']

    return txt


models_list = ["Azurro/APT3-1B-Instruct-v1"]

iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=2, placeholder="Wpisz tekst tutaj...", label="Wpisz tekst"),
        gr.Slider(value=500, label="Maksymalna długość", step=1, minimum=1, maximum=4000),
        gr.Slider(label="Temperatura", minimum=0.0, maximum=2.0, step=0.01, value=1.0),
        gr.Number(value=1.0, label="Top P", step=0.01, minimum=0.0, maximum=2.0)
    ],
    outputs=gr.Textbox(label="Wygenerowany tekst")
)

iface.launch()