Update main.py

main.py

@@ -1,13 +1,15 @@
+import math
 import torch
 import torch.nn as nn
-import math
-import re
 from fastapi import FastAPI
 from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 # ==========================================
-# 1. الثوابت (Constants)
+# 1. الإعدادات والثوابت (Constants)
 # ==========================================
+DEVICE = torch.device('cpu')  # إجبار العمل على المعالج العادي للسيرفر المجاني
+
 FATHA    = '\u064E'
 DAMMA    = '\u064F'
 KASRA    = '\u0650'
@@ -26,13 +28,12 @@ DIACRITIC_CLASSES = [
 ]
 IDX_TO_DIAC = {i: d for i, d in enumerate(DIACRITIC_CLASSES)}
 NUM_CLASSES = 15
-DEVICE = torch.device('cpu') # إجبار العمل على المعالج العادي للسيرفر المجاني
 
 def remove_diacritics(text):
     return ''.join(c for c in text if c not in ALL_DIACRITICS)
 
 # ==========================================
-# 2. بنية النموذج (Architecture Classes)
+# 2. بنية نموذج التشكيل (Architecture)
 # ==========================================
 class MultiHeadSelfAttention(nn.Module):
     def __init__(self, hidden_dim, num_heads=8, dropout=0.1):
@@ -116,18 +117,16 @@ class BiLSTMAttentionDiacritizer(nn.Module):
         return self.classifier(lstm_out)
 
 # ==========================================
-# 3. تحميل النموذج والأوزان
+# 3. تحميل النماذج (Model Loading)
 # ==========================================
-# تحميل الملف المرجعي وتوجيهه للـ CPU
-checkpoint = torch.load('Tashkeel_model.pt', map_location=DEVICE)
-
-# استخراج القاموس لمعرفة حجم المفردات
+print("جاري تحميل نموذج التشكيل...")
+# تأكد من أن اسم الملف هنا يطابق الملف المرفوع على Hugging Face
+CHECKPOINT_PATH = 'Tashkeel_model.pt' 
+checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
 char_to_idx = checkpoint['char_to_idx']
-VOCAB_SIZE = len(char_to_idx)
 
-# تهيئة النموذج بنفس الإعدادات التي تدرب عليها
 tashkeel_model = BiLSTMAttentionDiacritizer(
-    vocab_size=VOCAB_SIZE,
+    vocab_size=len(char_to_idx),
     embed_dim=256,
     hidden_dim=256,
     num_lstm_layers=3,
@@ -137,12 +136,20 @@ tashkeel_model = BiLSTMAttentionDiacritizer(
     dropout=0.3
 ).to(DEVICE)
 
-# تركيب الأوزان
 tashkeel_model.load_state_dict(checkpoint['model_state_dict'])
-tashkeel_model.eval() # تفعيل وضع الاستخدام
+tashkeel_model.eval()
+
+print("جاري تحميل نموذج AraBART...")
+# في حال كان مسار النموذج على Hugging Face مختلفاً، يرجى تعديل السطر التالي
+ARABART_MODEL_NAME = "UBC-NLP/AraBART-Morph-GEC" 
+arabart_tokenizer = AutoTokenizer.from_pretrained(ARABART_MODEL_NAME)
+arabart_model = AutoModelForSeq2SeqLM.from_pretrained(ARABART_MODEL_NAME)
+gec_pipeline = pipeline("text2text-generation", model=arabart_model, tokenizer=arabart_tokenizer)
+
+print("✅ تمت تهيئة جميع النماذج بنجاح!")
 
 # ==========================================
-# 4. دوال التشكيل والـ API
+# 4. دالة التشكيل المساعدة
 # ==========================================
 def diacritize_text(text, model, char_to_idx, max_len=200):
     clean = remove_diacritics(text)
@@ -180,16 +187,33 @@ def diacritize_text(text, model, char_to_idx, max_len=200):
         if d: result.append(d)
     return ''.join(result)
 
+# ==========================================
+# 5. إعداد الخادم ومسارات الـ API
+# ==========================================
 app = FastAPI()
 
 class TextRequest(BaseModel):
     text: str
 
+@app.get("/")
+def read_root():
+    return {
+        "status": "online",
+        "message": "مرحباً بك في واجهة برمجة تطبيقات نظام نحو (NAHW) للمعالجة الذكية للنصوص",
+        "endpoints": ["/tashkeel", "/spell-check", "/grammar-check"]
+    }
+
 @app.post("/tashkeel")
 def process_tashkeel(request: TextRequest):
     result = diacritize_text(request.text, tashkeel_model, char_to_idx)
-    return {"diacritized_text": result}
+    return {"result_text": result}
 
-@app.get("/")
-def read_root():
-    return {"message": "API نحو للتشكيل الآلي يعمل بنجاح!"}
+@app.post("/spell-check")
+def spell_check(request: TextRequest):
+    result = gec_pipeline(request.text)
+    return {"result_text": result[0]['generated_text']}
+
+@app.post("/grammar-check")
+def grammar_check(request: TextRequest):
+    result = gec_pipeline(request.text)
+    return {"result_text": result[0]['generated_text']}