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

title = "פייקיפדיה"

DESCRIPTION = """\
# צרו ערכים מטופשים 
מבוסס על [מודל דו-לשוני](https://huggingface.co/Norod78/SmolLM-135M-FakyPedia-EngHeb) התומך בעברית ובאנגלית המייצר טקסט מופרך בסגנון ויקיפדיה.  
לחצו על כפתור "שלח" כדי ליצור ערך "ויקיפדיה" מטופש ו/או מזויף בהתבסס על הכותרת שהזנתם
"""

article = """\
    מודל זה הרחיב את הטוקנייזר של  
    [SmolLM-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM-135M-Instruct)  
      המודל כּוּיַּיל לפייקיפדיה ע״י [דורון אדלר](https://linktr.ee/Norod78)  
        
"""

CUDA_AVAILABLE = torch.cuda.is_available()
device = torch.device("cuda" if CUDA_AVAILABLE else "cpu")

#model_id = "../SmolLM-135M-FakyPedia-EngHeb"
model_id = "Norod78/SmolLM-135M-FakyPedia-EngHeb"
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token_id = tokenizer.eos_token_id
bos_token = tokenizer.bos_token
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
model.generation_config.pad_token_id = tokenizer.pad_token_id

torch.manual_seed(1234)

@spaces.GPU
def generate_fakypedia(article_title: str):
    with torch.no_grad():
        result = ""
        string_to_tokenize= f"{bos_token}\\%{article_title}"
        input_ids = tokenizer(string_to_tokenize, return_tensors="pt").input_ids.to(device)        
        sample_outputs = model.generate(input_ids, do_sample=True,repetition_penalty=1.05, top_k = 40, top_p = 0.950, temperature=0.80, max_length=192, num_return_sequences=3)
        #sample_outputs = model.generate(input_ids, do_sample=True,repetition_penalty=1.2, temperature=0.50, max_length=192, num_return_sequences=3)
        if article_title == None or len(article_title) == 0:
            result += f"# תוצאות פייקיפדיה עם כותרות אקראיות  \n"
            article_title = ""
        else:
            result += f"# תוצאות פייקיפדיה עבור \"{article_title}\"  \n"
        for i, sample_output in enumerate(sample_outputs):
            decoded_output = tokenizer.decode(sample_output, skip_special_tokens=True)
            decoded_output = decoded_output.replace(f"\%{article_title}", f"##  {i+1}. {article_title}").replace("\%", " ").replace("\\n", "  \n")
            decoded_output = decoded_output.replace("##   \n", "\n")
            result += "{}\n".format(decoded_output)
        return result

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

demo.queue()
demo.launch()