customization

app.py

@@ -4,7 +4,7 @@ import torch
 from transformers import (
     Wav2Vec2ForCTC, Wav2Vec2Processor, 
     MarianMTModel, MarianTokenizer, 
-    BertForSequenceClassification, AutoTokenizer, AutoModel
+    BertForSequenceClassification, AutoModel, AutoTokenizer
 )
 
 # Detect device
@@ -24,12 +24,13 @@ translation_model = MarianMTModel.from_pretrained(translation_model_name).to(dev
 # AraBERT for Darija topic classification
 arabert_model_name = "aubmindlab/bert-base-arabert"
 arabert_tokenizer = AutoTokenizer.from_pretrained(arabert_model_name)
-arabert_model = BertForSequenceClassification.from_pretrained(arabert_model_name, num_labels=11).to(device)  # Adjusted to 11 labels for Darija
+arabert_model = BertForSequenceClassification.from_pretrained(arabert_model_name, num_labels=2).to(device)
+
 
 # BERT for English topic classification
 bert_model_name = "bert-base-uncased"
 bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
-bert_model = BertForSequenceClassification.from_pretrained(bert_model_name, num_labels=11).to(device)  # Adjusted to 11 labels for English
+bert_model = BertForSequenceClassification.from_pretrained(bert_model_name, num_labels=3).to(device)  
 
 # Libellés en Darija (Arabe et Latin)
 darija_topic_labels = [
@@ -61,6 +62,31 @@ english_topic_labels = [
     "Other"
 ]
 
+# New Function to Classify Topics by Keywords
+def classify_topic_by_keywords(text, topic_labels):
+    # Dictionnaire de mots-clés pour chaque topic
+    keywords = {
+        "خدمة العملاء": ["خدمة", "استفسار", "مساعدة", "دعم", "سؤال", "استفسار"],
+        "خدمة الاحتفاظ": ["احتفاظ", "تجديد", "خصم", "عرض", "العرض"],
+        "مشكلة في الفاتورة": ["فاتورة", "دفع", "مشكلة", "خطأ", "مبلغ"]
+    }
+    
+    # Convertir le texte en minuscule pour éviter les incohérences
+    text = text.lower()
+    
+    # Vérification de la présence des mots-clés dans le texte
+    topic_scores = {label: 0 for label in topic_labels}  # Initialiser le score des topics
+
+    for topic, words in keywords.items():
+        for word in words:
+            if word in text:
+                topic_scores[topic] += 1  # Incrémenter le score pour chaque mot trouvé
+    
+    # Retourner le topic avec le score le plus élevé
+    best_topic = max(topic_scores, key=topic_scores.get)
+    return best_topic
+
+
 def transcribe_audio(audio):
     """Convert audio to text, translate it, and classify topics in both Darija and English."""
     try:
@@ -77,14 +103,18 @@ def transcribe_audio(audio):
         # Translate to English
         translation = translate_text(transcription)
 
-        # Classify topics
+        # Classify topics using BERT models
         darija_topic = classify_topic(transcription, arabert_tokenizer, arabert_model, darija_topic_labels)
         english_topic = classify_topic(translation, bert_tokenizer, bert_model, english_topic_labels)
 
-        return transcription, translation, darija_topic, english_topic
+        # Classify topics using keywords-based classification
+        darija_keyword_topic = classify_topic_by_keywords(transcription, darija_topic_labels)
+        english_keyword_topic = classify_topic_by_keywords(translation, english_topic_labels)
+
+        return transcription, translation, darija_topic, english_topic, darija_keyword_topic, english_keyword_topic
 
     except Exception as e:
-        return f"Error processing audio: {str(e)}", "", "", ""
+        return f"Error processing audio: {str(e)}", "", "", "", "", ""
 
 def translate_text(text):
     """Translate Arabic text to English."""
@@ -111,11 +141,15 @@ with gr.Blocks() as demo:
 
     transcription_output = gr.Textbox(label="Transcription (Darija)")
     translation_output = gr.Textbox(label="Translation (English)")
-    darija_topic_output = gr.Textbox(label="Darija Topic Classification")
-    english_topic_output = gr.Textbox(label="English Topic Classification")
+    darija_topic_output = gr.Textbox(label="Darija Topic Classification (BERT)")
+    english_topic_output = gr.Textbox(label="English Topic Classification (BERT)")
+    darija_keyword_topic_output = gr.Textbox(label="Darija Topic Classification (Keywords)")
+    english_keyword_topic_output = gr.Textbox(label="English Topic Classification (Keywords)")
 
     submit_button.click(transcribe_audio, 
                         inputs=[audio_input], 
-                        outputs=[transcription_output, translation_output, darija_topic_output, english_topic_output])
+                        outputs=[transcription_output, translation_output, 
+                                 darija_topic_output, english_topic_output, 
+                                 darija_keyword_topic_output, english_keyword_topic_output])
 
 demo.launch()