Add Moroccan Darija extraction app5

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, pipeline
-import soundfile as sf
+from transformers import BertTokenizer, BertForSequenceClassification
 import librosa
 
 # Load models
@@ -10,42 +10,63 @@ processor = Wav2Vec2Processor.from_pretrained("boumehdi/wav2vec2-large-xlsr-moro
 transcription_model = Wav2Vec2ForCTC.from_pretrained("boumehdi/wav2vec2-large-xlsr-moroccan-darija")
 
 # Summarization model
-summarizer = pipeline("summarization", model="t5-small")
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+
+# Topic Classification Model (BERT for example)
+topic_model = BertForSequenceClassification.from_pretrained("your-finetuned-topic-model")
+topic_tokenizer = BertTokenizer.from_pretrained("your-finetuned-topic-model")
+
+# Function to resample audio to 16kHz if necessary
+def resample_audio(audio_path, target_sr=16000):
+    audio_input, original_sr = librosa.load(audio_path, sr=None)  # Load audio with original sampling rate
+    if original_sr != target_sr:
+        audio_input = librosa.resample(audio_input, orig_sr=original_sr, target_sr=target_sr)  # Resample to 16kHz
+    return audio_input, target_sr
 
 # Function to transcribe audio using Wav2Vec2
 def transcribe_audio(audio_path):
-    # Load and resample audio to 16kHz
-    audio_input, original_sample_rate = sf.read(audio_path)
-    if original_sample_rate != 16000:
-        audio_input = librosa.resample(audio_input, orig_sr=original_sample_rate, target_sr=16000)
+    # Load and preprocess audio
+    audio_input, sample_rate = resample_audio(audio_path)
+    inputs = processor(audio_input, sampling_rate=sample_rate, return_tensors="pt", padding=True)
 
-    # Process audio for transcription
-    inputs = processor(audio_input, sampling_rate=16000, return_tensors="pt", padding=True)
+    # Get predictions
     with torch.no_grad():
         logits = transcription_model(**inputs).logits
+
+    # Decode predictions
     predicted_ids = torch.argmax(logits, dim=-1)
     transcription = processor.batch_decode(predicted_ids)[0]
     return transcription
 
-# Function to analyze topics
-def analyze_topics(summary):
-    if "customer service" in summary.lower():
+# Function to classify the transcription into topics
+def classify_topic(transcription):
+    # Tokenize the transcription and pass it through the BERT classifier
+    inputs = topic_tokenizer(transcription, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    with torch.no_grad():
+        outputs = topic_model(**inputs)
+    
+    # Get the predicted label (0 for Customer Service, 1 for Retention Service, etc.)
+    predicted_class = torch.argmax(outputs.logits, dim=1).item()
+
+    # Map prediction to a topic
+    if predicted_class == 0:
         return "Customer Service"
-    elif "retention" in summary.lower():
-        return "Retention"
+    elif predicted_class == 1:
+        return "Retention Service"
     else:
         return "Other"
 
-# Function to transcribe, summarize, and analyze
-def transcribe_summarize_analyze(audio_file):
+# Function to transcribe, summarize, and classify topic
+def transcribe_and_summarize(audio_file):
     # Transcription
     transcription = transcribe_audio(audio_file)
 
     # Summarization
     summary = summarizer(transcription, max_length=50, min_length=10, do_sample=False)[0]["summary_text"]
 
-    # Topic Analysis
-    topic = analyze_topics(summary)
+    # Topic classification
+    topic = classify_topic(transcription)
+
     return transcription, summary, topic
 
 # Gradio Interface
@@ -57,11 +78,11 @@ outputs = [
 ]
 
 app = gr.Interface(
-    fn=transcribe_summarize_analyze,
+    fn=transcribe_and_summarize,
     inputs=inputs,
     outputs=outputs,
-    title="Moroccan Darija Audio Processing",
-    description="Upload an audio file in Moroccan Darija to get its transcription, a summarized version of the content, and an identified topic (e.g., Customer Service or Retention)."
+    title="Moroccan Darija Audio Transcription, Summarization, and Topic Classification",
+    description="Upload an audio file in Moroccan Darija to get its transcription, a summarized version of the content, and the detected topic."
 )
 
 # Launch the app