second

app.py

@@ -4,7 +4,7 @@ import torch
 from transformers import (
     Wav2Vec2ForCTC, Wav2Vec2Processor, 
     MarianMTModel, MarianTokenizer, 
-    BertForSequenceClassification, AutoModel, AutoTokenizer
+    BertForSequenceClassification, AutoTokenizer, AutoModel
 )
 
 # Detect device
@@ -24,13 +24,12 @@ 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=2).to(device)
-
+arabert_model = BertForSequenceClassification.from_pretrained(arabert_model_name, num_labels=11).to(device)  # Adjusted to 11 labels for Darija
 
 # 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).to(device)  
+bert_model = BertForSequenceClassification.from_pretrained(bert_model_name, num_labels=11).to(device)  # Adjusted to 11 labels for English
 
 # Libellés en Darija (Arabe et Latin)
 darija_topic_labels = [
@@ -62,7 +61,6 @@ english_topic_labels = [
     "Other"
 ]
 
-
 def transcribe_audio(audio):
     """Convert audio to text, translate it, and classify topics in both Darija and English."""
     try: