Update STT/sst.py

STT/sst.py

@@ -2,30 +2,30 @@ from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 import torchaudio
 import torch
 
-# تحميل المعالج والنموذج
-processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
-model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
+# تحميل المعالج والموديل العربي
+processor = Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-arabic")
+model = Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-arabic")
 
 def speech_to_text(audio_path):
     if audio_path is None:
-        raise ValueError("Audio path is None. Did you upload a file?")
-    
-    # تحميل الصوت
-    waveform, sampling_rate = torchaudio.load(audio_path)
+        raise ValueError("الصوت غير موجود")
 
-    # إذا كان ستيريو نخليه mono
+    # تحميل الملف الصوتي
+    waveform, sample_rate = torchaudio.load(audio_path)
+
+    # إذا الصوت ستيريو نحوله لمونو
     if waveform.shape[0] > 1:
-        waveform = waveform.mean(dim=0)
+        waveform = waveform.mean(dim=0).unsqueeze(0)
 
-    # إعادة تشكيل الصوت إذا كان غير 16kHz
-    if sampling_rate != 16000:
-        resampler = torchaudio.transforms.Resample(orig_freq=sampling_rate, new_freq=16000)
+    # إعادة تحويل التردد إلى 16000 لو كان مختلف
+    if sample_rate != 16000:
+        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
         waveform = resampler(waveform)
-    
-    # تجهيز البيانات للنموذج
+
+    # تجهيز الإدخال للنموذج
     input_values = processor(waveform.squeeze().numpy(), return_tensors="pt", sampling_rate=16000).input_values
 
-    # استنتاج الـ logits والتنبؤ
+    # تمرير البيانات للنموذج والحصول على النتائج
     with torch.no_grad():
         logits = model(input_values).logits
 
@@ -35,4 +35,3 @@ def speech_to_text(audio_path):
     transcription = processor.batch_decode(predicted_ids)
 
     return transcription[0]
-