Update goai_stt.py

goai_stt.py

@@ -5,18 +5,14 @@ from transformers import set_seed, Wav2Vec2ForCTC, AutoProcessor
 
 device = 0 if torch.cuda.is_available() else "cpu"
 
-
 def goai_stt(fichier):
     """
     Transcrire un fichier audio donné.
     
     Paramètres
     ---------- 
-    fichier: str
-        Le chemin d'accès au fichier audio.
-
-    device: str
-        GPU ou CPU
+    fichier: str | np.ndarray
+        Le chemin d'accès au fichier audio ou le tableau numpy.
         
     Return
     ---------- 
@@ -24,9 +20,8 @@ def goai_stt(fichier):
         Le texte transcrit.
     """
     
-    
     ### assurer reproducibilité
-    set_seed(2024) 
+    set_seed(2024)
     
     start_time = time.time()
     
@@ -34,11 +29,16 @@ def goai_stt(fichier):
     model_id = "anyantudre/wav2vec2-large-mms-1b-mos-V1"
 
     processor = AutoProcessor.from_pretrained(model_id)
-    model = Wav2Vec2ForCTC.from_pretrained(model_id, target_lang="mos", ignore_mismatched_sizes=True)
+    model = Wav2Vec2ForCTC.from_pretrained(model_id, target_lang="mos", ignore_mismatched_sizes=True).to(device)
+
+    if isinstance(fichier, str):
+        ### preprocessing de l'audio à partir d'un fichier
+        signal, sampling_rate = librosa.load(fichier, sr=16000)
+    else:
+        ### preprocessing de l'audio à partir d'un tableau numpy
+        signal, sampling_rate = fichier
 
-    ### preprocessing de l'audio
-    signal, sampling_rate =  librosa.load(fichier, sr=16000)
-    inputs = processor(signal, sampling_rate=16_000, return_tensors="pt", padding=True)
+    inputs = processor(signal, sampling_rate=16_000, return_tensors="pt", padding=True).to(device)
     
     ### faire l'inference
     with torch.no_grad():
@@ -48,4 +48,4 @@ def goai_stt(fichier):
     transcription = processor.decode(pred_ids)
 
     print("Temps écoulé: ", int(time.time() - start_time), " secondes")
-    return transcription
+    return transcription