import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

# 1. استخدام مستودع مفتوح ومفتوح الصلاحيات للـ Tokenizer لتفادي خطأ الـ Gated Repo
MODEL_ID = "unsloth/gemma-2b-it"
GGUF_REPO = "rof7iv/gemma"
GGUF_FILE = "gemma-4-E4B-it-ultra-uncensored-heretic-Q4_K_M.gguf"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# 2. تحميل الموديل باستخدام GGUF
model = AutoModelForCausalLM.from_pretrained(
    GGUF_REPO,
    gguf_file=GGUF_FILE,
    device_map="cpu"
)

def chat_with_model(message, history):
    conversation = []
    for user_prompt, bot_response in history:
        conversation.append({"role": "user", "content": user_prompt})
        conversation.append({"role": "assistant", "content": bot_response})
    conversation.append({"role": "user", "content": message})
    
    prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=512,
        temperature=0.7,
    )
    
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()
    
    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        yield partial_message

demo = gr.ChatInterface(
    fn=chat_with_model,
    title="Gemma Uncensored Cloud Server",
    description="سيرفر سحابي يعمل لخدمة تطبيقاتك."
)

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