src/model.py

"""
model.py - بناء وتحميل نموذج Transformer صغير لتصنيف حركات XO.
"""

from transformers import BertConfig, BertForSequenceClassification


def build_model(tokenizer):
    """
    ينشئ نموذج BERT صغير جداً بـ 9 رؤوس تصنيف (الخانات 0-8).
    المُدخل: tokenizer (لمعرفة حجم المفردات ورمز padding).
    المُخرج: النموذج الجاهز.
    """
    config = BertConfig(
        vocab_size=len(tokenizer),
        hidden_size=64,
        num_hidden_layers=2,
        num_attention_heads=2,
        intermediate_size=128,
        max_position_embeddings=16,
        num_labels=9,
        pad_token_id=tokenizer.pad_token_id,
    )

    model = BertForSequenceClassification(config)
    print(f"✅ تم بناء النموذج ({sum(p.numel() for p in model.parameters()):,} معلمة)")
    return model


def load_model(model_path="./xo_model"):
    """
    يحمّل نموذجاً مدرباً مسبقاً.
    """
    model = BertForSequenceClassification.from_pretrained(model_path)
    print(f"📂 تم تحميل النموذج من {model_path}")
    return model