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

title = "מחולל שירים"

DESCRIPTION = """\
# צרו שירים מטופשים 
המודל הוא [פיינטיון של גוגל גאממא 2 - 2ב׳](https://huggingface.co/Norod78/hebrew_lyrics-gemma2_2b-unsloth-gguf)  
כתבו פרומפט בסגנון ״כתוב לי בבקשה שיר על / המתאר / שמדבר על ____״
"""

article = """\
      המודל כּוּיַּיל ע״י [דורון אדלר](https://linktr.ee/Norod78)  
"""

#model_id = "./hebrew_lyrics-gemma2_2b"
model_id = "Norod78/hebrew_lyrics-gemma2_2b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16)

# model_id = "Norod78/hebrew_lyrics-gemma2_2b-unsloth-gguf"
# gguf_file_name = "hebrew_lyrics-gemma2_2b-unsloth.BF16.gguf"
# tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=gguf_file_name)
# model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, gguf_file=gguf_file_name).to("cpu")

torch.manual_seed(1234)

@spaces.GPU
def generate_song(prompt_text = ''):
    with torch.no_grad():
        result = ""
        input_template = tokenizer.apply_chat_template([{"role": "user", "content": prompt_text}], tokenize=False, add_generation_prompt=True)
        input_ids = tokenizer(input_template, return_tensors="pt").to(model.device)
        #sample_outputs = model.generate(**input_ids, max_new_tokens=512 , repetition_penalty=1.1, temperature=0.4, top_p=0.95, top_k=40, do_sample = True)
        sample_outputs = model.generate(**input_ids, max_new_tokens=384 , repetition_penalty=1.1, temperature=0.6, top_p=0.4, top_k=40, do_sample = True)
        #sample_outputs = model.generate(**input_ids, max_new_tokens=512 , repetition_penalty=1.1, temperature=0.5, do_sample = True)
        decoded_output = tokenizer.batch_decode(sample_outputs, skip_special_tokens=True)[0]
        result = decoded_output.replace("user\n", "משתמש:\n").replace("model\n", "\nמודל:\n")
        return result



demo = gr.Interface(
    generate_song,    
    inputs=gr.Textbox(lines=1, label="בקשו שיר", rtl=True),
    outputs=gr.Textbox(label="הפלט של המודל", rtl=True),
    title=title,
    description=DESCRIPTION,
    article=article,
    examples=["תנו לשמש לעלות, לבוקר להאיר", "כתוב לי בבקשה שיר על תפוח אדמה עם חרדה חברתית", "שיר המתאר את חייהם של חזירים בצבע סגול עם כנפיים וההרפתקאות המעופפות שלהם", "שיר על פתח תקווה", "שיר על חדי קרן וורודים 🦄"],
    allow_flagging="never",
)

demo.queue()
demo.launch()