File size: 479 Bytes
1784a22
c1f16ee
99757c1
 
 
 
 
 
 
 
 
c1f16ee
 
1784a22
 
99757c1
 
c1f16ee
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import pickle

from transformers import AutoTokenizer


def classify_arabic_dialect(text:str, model, tokenizer) -> str:
    text_embeddings = tokenizer(text, padding=True)
    predicted_class = model.predict(text_embeddings)

    return predicted_class


def main():
    with open("../models/logistic_regression.pkl", "rb") as f:
        model = pickle.load(f)

    tokenizer = AutoTokenizer.from_pretrained("moussaKam/AraBART")
    return


if __name__ == "__main__":
    main()