File size: 914 Bytes
ec4a7b0
1784a22
c1f16ee
7ecfa8c
9a23b5c
 
45c84cc
99757c1
 
ec4a7b0
 
 
 
 
 
 
 
99757c1
9a23b5c
 
 
99757c1
7ecfa8c
 
 
 
 
c1f16ee
7ecfa8c
 
 
 
 
c1f16ee
 
 
7ecfa8c
 
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
import os
import pickle

import gradio as gr
from transformers import AutoModel, AutoTokenizer

from .utils import extract_hidden_state


models_dir = os.path.join(os.path.dirname(__file__), '..', 'models')
model_file = os.path.join(models_dir, 'logistic_regression.pkl')

if os.path.exists(model_file):
    with open(model_file, "rb") as f:
        model = pickle.load(f)
else:
    print(f"Error: {model_file} not found.")

model_name = "moussaKam/AraBART"
tokenizer = AutoTokenizer.from_pretrained(model_name)
language_model = AutoModel.from_pretrained(model_name)

def classify_arabic_dialect(text):
    text_embeddings = extract_hidden_state(text, tokenizer, language_model)
    predicted_class = model.predict(text_embeddings)[0]
    
    return predicted_class

demo = gr.Interface(
    fn=classify_arabic_dialect,
    inputs=["text"],
    outputs=["text"],
)


if __name__ == "__main__":
    demo.launch()