ZiadElshayeb/Arabic-Dialect-Classification

1import gradio as gr
2from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
3import tensorflow as tf
4
5# Load Model and Tokenizer
6model_path = "ZiadElshayeb/dialect-classification-model"
7
8def load_model():
9    tokenizer = AutoTokenizer.from_pretrained(model_path)
10    model = TFAutoModelForSequenceClassification.from_pretrained(model_path)
11    return tokenizer, model
12
13tokenizer, model = load_model()
14
15# Label Mapping
16label_map = {0: "Egyptian", 1: "Gulf", 2: "Levant", 3: "Maghrebi", 4: "Southern & Mesopotamian"}
17
18# Function to Predict Dialect
19def predict_dialect(text):
20    # Tokenize input text
21    inputs = tokenizer(text, return_tensors="tf", padding=True, truncation=True, max_length=128)
22
23    # Run Model Prediction
24    predictions = model(inputs)
25    predicted_class_id = tf.math.argmax(predictions.logits, axis=1).numpy()[0]
26
27    # Get Dialect Label
28    predicted_class = label_map.get(predicted_class_id, "Unknown")
29    return f"Predicted Dialect: {predicted_class}"
30
31# Create Gradio Interface
32interface = gr.Interface(
33    fn=predict_dialect,
34    inputs=gr.Textbox(lines=3, placeholder="Enter Arabic text..."),
35    outputs="text",
36    title="Arabic Dialects Classification",
37    description="Enter Arabic text and the model will classify it into one of five dialects.",
38)
39
40# Launch the Gradio App
41if __name__ == "__main__":
42    interface.launch(debug=True)
43