Update app.py

app.py

@@ -1,10 +1,11 @@
 import torch
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, T5ForConditionalGeneration, T5Tokenizer
+from transformers import pipeline, AutoTokenizer, M2M100ForConditionalGeneration
+from tokenization_small100 import SMALL100Tokenizer
 import numpy as np
 from pydub import AudioSegment
 
-# Load the pipeline for speech recognition and translation
+# Load the pipeline for speech recognition
 pipe = pipeline(
     "automatic-speech-recognition",
     model="DrishtiSharma/whisper-large-v2-hausa",
@@ -12,9 +13,9 @@ pipe = pipeline(
 )
 
 # Load the new translation model and tokenizer
-model_name = 'jbochi/madlad400-3b-mt'
-model = T5ForConditionalGeneration.from_pretrained(model_name)
-tokenizer = T5Tokenizer.from_pretrained(model_name)
+model_name = 'alirezamsh/small100'
+model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+tokenizer = SMALL100Tokenizer.from_pretrained(model_name)
 
 tts = pipeline("text-to-speech", model="Baghdad99/english_voice_tts")
 
@@ -44,22 +45,11 @@ def translate_speech(audio_file):
         print("The output does not contain 'text'")
         return
 
-    # Use the translation pipeline to translate the transcription
-    translated_text = translator(transcription, return_tensors="pt")
-    print(f"Translated text: {translated_text}")  # Print the translated text to see what it contains
-
-    # Check if the translated text contains 'generated_token_ids'
-    if 'generated_token_ids' in translated_text[0]:
-        # Decode the tokens into text
-        translated_text_str = translator.tokenizer.decode(translated_text[0]['generated_token_ids'])
-    else:
-        print("The translated text does not contain 'generated_token_ids'")
-        return
-
-   # Use the new translation model to translate the transcription
+    # Use the new translation model to translate the transcription
     text = "translate Hausa to English: " + transcription
-    input_ids = tokenizer.encode(text, return_tensors="pt")
-    outputs = model.generate(input_ids=input_ids)
+    tokenizer.tgt_lang = "en"
+    encoded_text = tokenizer(text, return_tensors="pt")
+    outputs = model.generate(**encoded_text)
 
     # Decode the tokens into text
     translated_text_str = tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -93,4 +83,3 @@ iface = gr.Interface(
 )
 
 iface.launch()
-