import os
from threading import Thread
import gradio as gr
import torch
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TextIteratorStreamer,
StoppingCriteria,
StoppingCriteriaList,
)
from openai import OpenAI
# ==========================================
# النماذج المتاحة (مساران):
# 1) Qwen3.5-2B → عبر Transformers، يدعم وضع التفكير (thinking) القابل للتبديل.
# 2) GPT-5 → عبر الـ OpenAI API.
# النموذج المحلي يُحمّل كسولاً عند أول استخدام لتوفير الذاكرة عند الإقلاع.
# ==========================================
QWEN_MODEL = "Qwen/Qwen3.5-2B"
GPT_LABEL = "GPT-5 (OpenAI API)"
# حدود الرموز (Tokens):
# - رسالة المستخدم الواحدة: 512 رمزاً كحد أقصى (يُقصّ الزائد).
# - كامل سياق المحادثة المُرسل للنموذج: 4096 رمزاً (نُبقي الأحدث).
MAX_MSG_TOKENS = 1024
MAX_CHAT_TOKENS = 4096
# ميزانيتان منفصلتان للتوليد في وضع التفكير:
# - THINK_BUDGET : أقصى عدد رموز يُسمح به لكواليس التفكير. إن نفدت دون أن يُغلق
# النموذج التفكير () نُغلقه قسراً وننتقل لتوليد الجواب.
# - ANSWER_BUDGET: ميزانية مستقلة للجواب، فلا يبتلع التفكيرُ حصة الجواب أبداً.
# بهذا نضمن وصول جواب فعلي حتى لو أطال النموذج في التفكير.
THINK_BUDGET =512
ANSWER_BUDGET =512
THINK_END_ID = 248069
_qwen_models = {} # سجلّ نماذج Qwen المحمّلة عبر Transformers
# ------------------------------------------
# تحميل Qwen عبر Transformers (bfloat16 على الـ CPU لتقليل ذروة الذاكرة)
# ------------------------------------------
def get_qwen_model(model_id):
if model_id not in _qwen_models:
print(f"⏳ جاري تحميل نموذج {model_id} ...")
tok = AutoTokenizer.from_pretrained(model_id)
if torch.cuda.is_available():
# على الـ GPU: تحميل عادي بالدقة التلقائية
mdl = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
else:
# على الـ CPU: نحمّل بدقة bfloat16 بدون تكميم int8.
# السبب: تكميم float32→int8 يُحدث ذروة ذاكرة ضخمة (~10-11GB لنموذج 2B) تسبب OOM
# على الـ Space المجاني. تحميل bfloat16 يُبقي الذروة عند ~نصف ذلك فتتّسع النماذج الأكبر،
# مقابل توليد أبطأ قليلاً على الـ CPU (لا تسريع int8).
print("🔻 تحميل بدقة bfloat16 على CPU لتقليل ذروة استهلاك الذاكرة...")
mdl = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
_qwen_models[model_id] = (tok, mdl)
return _qwen_models[model_id]
# ==========================================
# توليد Qwen عبر Transformers مع دعم وضع التفكير القابل للتبديل
# Qwen3 يوصي: وضع التفكير temp=0.6/top_p=0.95 — الوضع العادي temp=0.7/top_p=0.8 — وtop_k=20 دائماً
# ==========================================
def _truncate_text_tokens(tok, text, max_tokens):
"""يقصّ النص إلى max_tokens رمزاً كحد أقصى (نُبقي البداية).
نتجاهل ما ليس نصاً (None/غير str) تماماً كما يفعل قالب المحادثة، تفادياً لخطأ المُرمِّز."""
if not isinstance(text, str) or not text:
return text
ids = tok(text, add_special_tokens=False).input_ids
if len(ids) <= max_tokens:
return text
return tok.decode(ids[:max_tokens], skip_special_tokens=True)
def qwen_generate(messages, enable_thinking, tok, mdl, max_new_tokens=1024):
temperature = 0.6 if enable_thinking else 0.7
top_p = 0.95 if enable_thinking else 0.8
device = next(mdl.parameters()).device
# حدّ رسالة المستخدم: نقصّ كل رسالة مستخدم إلى MAX_MSG_TOKENS رمزاً قبل بناء البرومبت
messages = [dict(m) for m in messages]
for m in messages:
if m["role"] == "user":
m["content"] = _truncate_text_tokens(tok, m["content"], MAX_MSG_TOKENS)
text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking)
model_inputs = tok([text], return_tensors="pt").to(device)
# حدّ سياق المحادثة: إن تجاوز الإدخال MAX_CHAT_TOKENS نُبقي آخر (الأحدث) رموزه فقط
input_ids = model_inputs.input_ids
attention_mask = model_inputs.attention_mask
if input_ids.shape[-1] > MAX_CHAT_TOKENS:
input_ids = input_ids[:, -MAX_CHAT_TOKENS:]
attention_mask = attention_mask[:, -MAX_CHAT_TOKENS:]
generated_ids = mdl.generate(
input_ids=input_ids,
attention_mask=attention_mask,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
top_p=top_p,
top_k=20,
)
return generated_ids[0][len(input_ids[0]):].tolist()
# ==========================================
# معيار إيقاف: نوقف التوليد فور إنتاج رمز نهاية التفكير
# ==========================================
class _StopOnToken(StoppingCriteria):
def __init__(self, token_id):
self.token_id = token_id
def __call__(self, input_ids, scores, **kwargs):
return input_ids[0, -1].item() == self.token_id
# ==========================================
# صياغة العرض: نغلّف التفكير في صندوق قابل للطي ثم نُلحق الجواب
# ==========================================
def _compose_output(badge, thinking, answer):
if thinking:
return f"{badge}🧠 كواليس تفكير النموذج
\n\n{thinking}\n\n