File size: 1,830 Bytes
45542e5 365abf6 45542e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
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()
# 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))
|