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

model_id = "my0919175/Sovythos-66M-Base"

# 1. تحميل التوكنايزر والموديل مع الموافقة التلقائية
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    use_safetensors=True
)

eos_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
model.config.eos_token_id = eos_id
model.config.pad_token_id = eos_id

# 2. كلاس الحماية لمنع الهلوسة والـ NaN والـ Inf
class AntiNanLogitsProcessor(LogitsProcessor):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        scores[torch.isnan(scores)] = 0.0
        scores[torch.isinf(scores)] = torch.where(scores[torch.isinf(scores)] > 0, 10000.0, -10000.0)
        return scores

logits_processor_list = LogitsProcessorList([AntiNanLogitsProcessor()])

# 3. دالة معالجة النصوص والتوليد المستمر (Streaming)
def respond(
    message,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    prompt = ""
    if system_message:
        prompt += f"<|system|>\n{system_message}<|endoftext|>\n"
    
    # قراءة التاريخ بمرونة تامة للتوافق مع كل إصدارات Gradio
    for turn in history:
        if isinstance(turn, dict):  # تنسيق Gradio 5
            if turn["role"] == "user":
                prompt += f"<|user|>\n{turn['content']}<|endoftext|>\n"
            elif turn["role"] == "assistant":
                prompt += f"<|assistant|>\n{turn['content']}<|endoftext|>\n"
        else:  # تنسيق Gradio 4 (Tuples/Lists)
            prompt += f"<|user|>\n{turn[0]}<|endoftext|>\n"
            if turn[1]:
                prompt += f"<|assistant|>\n{turn[1]}<|endoftext|>\n"
    
    prompt += f"<|user|>\n{message}<|endoftext|>\n<|assistant|>\n"
    
    inputs = tokenizer(prompt, return_tensors="pt")
    if "attention_mask" not in inputs:
        inputs["attention_mask"] = torch.ones_like(inputs["input_ids"])
        
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
    
    generation_kwargs = dict(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_new_tokens=max_tokens,
        eos_token_id=eos_id,
        pad_token_id=eos_id,
        do_sample=True if temperature > 0.05 else False,
        temperature=max(temperature, 1e-2),
        top_k=40,
        top_p=top_p,
        repetition_penalty=1.05,
        logits_processor=logits_processor_list,
        streamer=streamer,
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    response = ""
    for new_text in streamer:
        response += new_text
        yield response

# 4. بناء واجهة التشات (تم حذف باراميتر type المتسبب في الكراش)
chatbot = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a helpful AI assistant.", label="System message"),
        gr.Slider(minimum=1, maximum=512, value=128, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.6, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.85,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ]
)

with gr.Blocks() as demo:
    chatbot.render()

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