replace main with tashkeel model

main.py

@@ -1,7 +1,57 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+#from shakkala import Shakkala
+from pathlib import Path
+import torch
+from eo_pl import TashkeelModel as TashkeelModelEO
+from ed_pl import TashkeelModel as TashkeelModelED
+from tashkeel_tokenizer import TashkeelTokenizer
+from utils import remove_non_arabic
 
 app = FastAPI()
 
-@app.get("/")
-def read_root():
-    return {"Hello": "World!"}
+# Load CaTT models
+tokenizer = TashkeelTokenizer()
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+max_seq_len = 1024    
+    
+eo_ckpt_path = Path(__file__).parent / 'models/best_eo_mlm_ns_epoch_193.pt'
+ed_ckpt_path = Path(__file__).parent / 'models/best_ed_mlm_ns_epoch_178.pt'
+    
+# Load Encoder-Only model
+eo_model = TashkeelModelEO(tokenizer, max_seq_len=max_seq_len, 
+                              n_layers=6, learnable_pos_emb=False)
+eo_model.load_state_dict(torch.load(eo_ckpt_path, map_location=device))
+eo_model.eval().to(device)
+    
+# Load Encoder-Decoder model
+ed_model = TashkeelModelED(tokenizer, max_seq_len=max_seq_len,
+                              n_layers=3, learnable_pos_emb=False)
+ed_model.load_state_dict(torch.load(ed_ckpt_path, map_location=device))
+ed_model.eval().to(device)
+
+
+class CattRequest(BaseModel):
+    text: str
+    model_type: str  # "Encoder-Only" or "Encoder-Decoder"
+
+
+    
+
+
+@app.post("/catt")
+def infer_catt(request: CattRequest):
+    try:
+        input_text = remove_non_arabic(request.text)
+        batch_size = 16
+        verbose = True
+        
+        if request.model_type == 'Encoder-Only':
+            output_text = eo_model.do_tashkeel_batch([input_text], batch_size, verbose)
+        else:
+            output_text = ed_model.do_tashkeel_batch([input_text], batch_size, verbose)
+            
+        return {"result": output_text[0]}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+