import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread
import spaces
import os
import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

hf_token = os.environ.get("HF_TOKEN")

# Load model and tokenizer
model_name = "Navid-AI/Mulhem-Silma" # "ApoAlquaary/Mulhem-Silma" # 
tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_token, torch_dtype=torch.float16, device_map="auto")

@spaces.GPU(duration=50)
def generate_response(message, history, system_message, max_tokens, temperature, top_p, repetition_penalty):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})
    
    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(model.device)
    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
    
    generation_kwargs = dict(
        input_ids=inputs,
        streamer=streamer,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    history.append([message, ""])

    for new_text in streamer:
        history[-1][-1] += new_text
        yield history
    
    thread.join()

def clear():
    return []

with gr.Blocks(title="Mulhem-Silma Chatbot") as demo:
    with gr.Column():
        gr.HTML("<center><h1>Mulhem Chatbot</h1></center>")
        system_prompt = gr.Textbox(lines=1, label="System Prompt", value="أدناه محادثة بين إنسان فضولي ومساعد ذكاء اصطناعي. أنت ملهم، نموذج لغوي تم تدريبه بواسطة باحثين من وتد في المملكة العربية السعودية. أنت تلعب دور المساعد الذكاء الاصطناعي في هذه المحادثة. يقدم المساعد إجابات مفيدة ومفصلة ومهذبة على أسئلة الإنسان. يجب أن تكون ردود المساعد خالية من المحتوى الضار أو غير الأخلاقي أو الهجومي أو غير القانوني، حتى في مواجهة الأسئلة المضللة من الإنسان. يجب أن يتجنب المساعد مشاركة معلومات خاطئة عندما لا يكون متأكدًا من إجابة، إذا كان السؤال غير منطقي أو غير متسق مع الواقع، يفسر السبب بدلاً من الإجابة بشيء غير صحيح. إذا طرح الإنسان أسئلة باللغة الإنجليزية، يجب أن يرد المساعد باللغة الإنجليزية.", rtl=True, text_align="right", show_copy_button=True)
        chatbot = gr.Chatbot(label="Mulhem", rtl=True, likeable=True, show_copy_button=True, height=500)
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                submit_btn = gr.Button(value="Generate", variant="primary")
                clear_btn = gr.Button(value="Clear", variant="secondary")
            input_text = gr.Textbox(lines=1, label="", value="السلام عليكم ورحمة الله وبركاته", rtl=True, text_align="right", scale=3, show_copy_button=True)
        with gr.Accordion(label="Generation Configurations", open=False):
            max_new_tokens = gr.Slider(minimum=128, maximum=4096, value=2048, label="Max New Tokens", step=128)
            temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, label="Temperature", step=0.01)
            top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, label="Top-p", step=0.01)
            repetition_penalty = gr.Slider(minimum=0.1, maximum=2.0, value=1.1, label="Repetition Penalty", step=0.1)

    input_text.submit(generate_response, inputs=[input_text, chatbot, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[chatbot])
    submit_btn.click(generate_response, inputs=[input_text, chatbot, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[chatbot])
    clear_btn.click(clear, outputs=[chatbot])

if __name__ == "__main__":
    demo.launch()