File size: 1,216 Bytes
ec4a7b0
1784a22
c1f16ee
ab54966
9a23b5c
 
ec4a7b0
99757c1
ab54966
99757c1
ec4a7b0
 
 
 
 
 
 
 
99757c1
9a23b5c
 
 
99757c1
c1f16ee
ab54966
 
 
 
 
 
 
6893bb5
9a23b5c
6893bb5
 
ab54966
 
4b51583
ab54966
 
 
 
c1f16ee
 
 
 
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
import os
import pickle

from flask import Flask, request, jsonify
from transformers import AutoModel, AutoTokenizer

from src.utils import extract_hidden_state

app = Flask(__name__)

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)


@app.route("/classify", methods=["POST"])
def classify_arabic_dialect():
    try:
        data = request.json
        text = data.get("text")
        if not text:
            return jsonify({"error": "No text has been received"}), 400
        
        text_embeddings = extract_hidden_state(text, tokenizer, language_model)
        predicted_class = model.predict(text_embeddings)[0]
        
        return jsonify({"class": predicted_class}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


def main():
    app.run(debug=True)


if __name__ == "__main__":
    main()