classificationV1

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import librosa
 import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, MarianMTModel, MarianTokenizer
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, MarianMTModel, MarianTokenizer, BertForSequenceClassification, AutoModel, AutoTokenizer
 
 # Charger le modèle de transcription pour le Darija
 model = Wav2Vec2ForCTC.from_pretrained("boumehdi/wav2vec2-large-xlsr-moroccan-darija")
@@ -12,23 +12,39 @@ translation_model_name = "Helsinki-NLP/opus-mt-ar-en"
 translation_model = MarianMTModel.from_pretrained(translation_model_name)
 translation_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
 
+
+
+# Load AraBERT for Darija topic classification
+arabert_model_name = "aubmindlab/bert-base-arabert"
+arabert_tokenizer = AutoTokenizer.from_pretrained(arabert_model_name)
+arabert_model = AutoModel.from_pretrained(arabert_model_name)
+
+# Load 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=3)  # Adjust labels as needed
+
+darija_topic_labels = ["Customer Service", "Retention Service", "Billing Issue"]  # Adjust for Darija topics
+english_topic_labels = ["Support Request", "Subscription Issue", "Payment Dispute"]  # Adjust for English topics
+
+
 def transcribe_audio(audio):
-    """Convertir l'audio en texte et le traduire en anglais"""
-    # Charger et prétraiter l'audio
+    """Convert audio to text, translate it, and classify topics in both Darija and English"""
     audio_array, sr = librosa.load(audio, sr=16000)
     input_values = processor(audio_array, return_tensors="pt", padding=True).input_values
 
-    # Obtenir les prédictions du modèle
     logits = model(input_values).logits
     tokens = torch.argmax(logits, axis=-1)
 
-    # Décoder la transcription en Darija
     transcription = processor.decode(tokens[0])
-
-    # Traduire en anglais
     translation = translate_text(transcription)
 
-    return transcription, translation
+    # Classify topics for both Darija and English
+    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
+
 
 def translate_text(text):
     """Traduire le texte de l'arabe vers l'anglais"""
@@ -37,15 +53,31 @@ def translate_text(text):
     translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
     return translated_text
 
+def classify_topic(text, tokenizer, model, topic_labels):
+    """Classify topic using BERT-based models"""
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predicted_class = torch.argmax(outputs.logits, dim=1).item()
+
+    return topic_labels[predicted_class] if predicted_class < len(topic_labels) else "Other"
+
+
 # Interface utilisateur avec Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎙️ Speech-to-Text & Translation")
+    gr.Markdown("# 🎙️ Speech-to-Text, Translation & Topic Classification")
     
     audio_input = gr.Audio(type="filepath", label="Upload Audio or Record")
-    submit_button = gr.Button("Transcribe & Translate")
+    submit_button = gr.Button("Process")
+
     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")
 
-    submit_button.click(transcribe_audio, inputs=[audio_input], outputs=[transcription_output, translation_output])
+    submit_button.click(transcribe_audio, 
+                        inputs=[audio_input], 
+                        outputs=[transcription_output, translation_output, darija_topic_output, english_topic_output])
 
 demo.launch()
+