yo

app.py

@@ -15,51 +15,38 @@ logging.basicConfig(
     level=logging.DEBUG,
     format='%(asctime)s - %(levelname)s - %(message)s'
 )
-
-checkpoint_dir = "facebook/final_m2m100"
-# Initialize translation model
+# Update the model loading section
 try:
-    logging.info(f"Attempting to load M2M100 from {checkpoint_dir}")
+    # Try to load custom model
+    checkpoint_dir = "bishaltwr/final_m2m100"
+    logging.info(f"Attempting to load custom M2M100 from {checkpoint_dir}")
     tokenizer = M2M100Tokenizer.from_pretrained(checkpoint_dir)
-    logging.info("M2M100 tokenizer loaded successfully")
-    
     model_m2m = M2M100ForConditionalGeneration.from_pretrained(checkpoint_dir)
-    logging.info("M2M100 model loaded successfully")
-    
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    logging.info(f"Using device: {device}")
-    
-    model_m2m.to(device)
-    m2m_available = True
-    logging.info("M2M100 model ready for use")
+    logging.info("Custom M2M100 model loaded successfully")
 except Exception as e:
-    logging.error(f"Error loading M2M100 model: {e}")
-    m2m_available = False
-    logging.info("Setting m2m_available to False")
+    logging.error(f"Error loading custom M2M100 model: {e}")
+    # Fall back to official model
+    checkpoint_dir = "facebook/m2m100_418M"
+    logging.info(f"Attempting to load official M2M100 from {checkpoint_dir}")
+    tokenizer = M2M100Tokenizer.from_pretrained(checkpoint_dir)
+    model_m2m = M2M100ForConditionalGeneration.from_pretrained(checkpoint_dir)
+    logging.info("Official M2M100 model loaded successfully")
+
+# Set device after model loading
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logging.info(f"Using device: {device}")
+model_m2m.to(device)
     
 # Initialize ASR model
 model_id = "bishaltwr/wav2vec2-large-mms-1b-nepali"
-try:
-    processor = AutoProcessor.from_pretrained(model_id)
-    model_asr = Wav2Vec2ForCTC.from_pretrained(model_id, ignore_mismatched_sizes=True)
-    asr_available = True
-except Exception as e:
-    logging.error(f"Error loading ASR model: {e}")
-    asr_available = False
+processor = AutoProcessor.from_pretrained(model_id)
+model_asr = Wav2Vec2ForCTC.from_pretrained(model_id, ignore_mismatched_sizes=True)
 
 # Initialize X-Transformer model
-try:
-    from inference import translate as xtranslate
-    xtransformer_available = True
-except Exception as e:
-    logging.error(f"Error loading XTransformer model: {e}")
-    xtransformer_available = False
+from inference import translate as xtranslate
 
 def m2m_translate(text, source_lang, target_lang):
     """Translation using M2M100 model"""
-    if not m2m_available:
-        return "M2M100 model not available"
-    
     tokenizer.src_lang = source_lang
     inputs = tokenizer(text, return_tensors="pt").to(device)
     translated_tokens = model_m2m.generate(
@@ -71,9 +58,6 @@ def m2m_translate(text, source_lang, target_lang):
 
 def transcribe_audio(audio_path, language="npi"):
     """Transcribe audio using ASR model"""
-    if not asr_available:
-        return "ASR model not available"
-    
     import librosa
     audio, sr = librosa.load(audio_path, sr=16000)
     processor.tokenizer.set_target_lang(language)
@@ -123,69 +107,55 @@ def translate_text(text, model_choice, source_lang=None, target_lang=None):
         target_lang = "ne" if source_lang == "en" else "en"
     
     # Choose the translation model
-    if model_choice == "XTransformer" and xtransformer_available:
+    if model_choice == "XTransformer":
         return xtranslate(text)
-    elif model_choice == "M2M100" and m2m_available:
+    elif model_choice == "M2M100":
         return m2m_translate(text, source_lang=source_lang, target_lang=target_lang)
     else:
         return "Selected model is not available"
 
 # Set up the Gradio interface
 with gr.Blocks(title="Nepali-English Translator") as demo:
-    gr.Markdown("# Nepali-English Translation Service")
+    gr.Markdown("# Nepali-English Translator")
     gr.Markdown("Translate between Nepali and English, transcribe audio, and convert text to speech.")
-    
-    # Set up tabs for different functions
-    with gr.Tabs():
-        # Text Translation Tab
-        with gr.TabItem("Text Translation"):
-            with gr.Row():
-                with gr.Column():
-                    text_input = gr.Textbox(label="Input Text", lines=5)
-                    
-                    with gr.Row():
-                        model_choice = gr.Radio(
-                            choices=["XTransformer", "M2M100"], 
-                            value="XTransformer", 
-                            label="Translation Model"
-                        )
-                        
-                    with gr.Row():
-                        source_lang = gr.Dropdown(
-                            choices=["Auto-detect", "en", "ne"], 
-                            value="Auto-detect", 
-                            label="Source Language",
-                            visible=True
-                        )
-                        target_lang = gr.Dropdown(
-                            choices=["Auto-select", "en", "ne"], 
-                            value="Auto-select", 
-                            label="Target Language",
-                            visible=True
-                        )
-                    
-                    translate_button = gr.Button("Translate")
-                
-                with gr.Column():
-                    translation_output = gr.Textbox(label="Translation Output", lines=5)
-                    tts_button = gr.Button("Convert to Speech")
-                    audio_output = gr.Audio(label="Audio Output")
+    gr.Markdown("Aakash Budhathoki, Apekshya Subedi, Bishal Tiwari, Kebin Malla. - Kantipur Engineering College.")
+    
+    with gr.Column():
+        gr.Markdown("### Speech to Text")
+        audio_input = gr.Audio(label="Upload or Record Audio", type="filepath")
+        asr_language = gr.Radio(
+            choices=["eng", "npi"], 
+            value="npi", 
+            label="Speech Language"
+        )
+        transcribe_button = gr.Button("Transcribe")
+        transcription_output = gr.Textbox(label="Transcription Output", lines=3)
+        
+        gr.Markdown("### Text Translation")
+        model_choice = gr.Dropdown(
+            choices=["XTransformer", "M2M100"], 
+            value="M2M100", 
+            label="Translation Model"
+        )
+        source_lang = gr.Dropdown(
+            choices=["Auto-detect", "en", "ne"], 
+            value="Auto-detect", 
+            label="Source Language"
+        )
+        target_lang = gr.Dropdown(
+            choices=["Auto-select", "en", "ne"], 
+            value="Auto-select", 
+            label="Target Language"
+        )
+        translate_button = gr.Button("Translate")
+        translation_output = gr.Textbox(label="Translation Output", lines=5)
         
-        # Speech to Text Tab
-        with gr.TabItem("Speech to Text"):
-            with gr.Column():
-                audio_input = gr.Audio(label="Upload or Record Audio", type="filepath")
-                asr_language = gr.Radio(
-                    choices=["eng", "npi"], 
-                    value="npi", 
-                    label="Speech Language"
-                )
-                transcribe_button = gr.Button("Transcribe")
-                transcription_output = gr.Textbox(label="Transcription Output", lines=3)
+        gr.Markdown("### Text to Speech")
+        tts_button = gr.Button("Convert to Speech")
+        audio_output = gr.Audio(label="Audio Output")
     
     # Define event handlers
     def process_translation(text, model, src_lang, tgt_lang):
-        logging.info(f"Processing translation: text={text}, model={model}, src_lang={src_lang}, tgt_lang={tgt_lang}")
         if src_lang == "Auto-detect":
             src_lang = None
         if tgt_lang == "Auto-select":
@@ -193,19 +163,23 @@ with gr.Blocks(title="Nepali-English Translator") as demo:
         return translate_text(text, model, src_lang, tgt_lang)
     
     def process_tts(text):
-        logging.info(f"Processing TTS: text={text}")
         return text_to_speech(text)
     
     def process_transcription(audio_path, language):
-        logging.info(f"Processing transcription: audio_path={audio_path}, language={language}")
         if not audio_path:
             return "Please upload or record audio"
         return transcribe_audio(audio_path, language)
     
     # Connect the components
+    transcribe_button.click(
+        process_transcription,
+        inputs=[audio_input, asr_language],
+        outputs=transcription_output
+    )
+    
     translate_button.click(
         process_translation, 
-        inputs=[text_input, model_choice, source_lang, target_lang], 
+        inputs=[transcription_output, model_choice, source_lang, target_lang], 
         outputs=translation_output
     )
     
@@ -214,40 +188,7 @@ with gr.Blocks(title="Nepali-English Translator") as demo:
         inputs=translation_output,
         outputs=audio_output
     )
-    
-    transcribe_button.click(
-        process_transcription,
-        inputs=[audio_input, asr_language],
-        outputs=transcription_output
-    )
-
-    # Explicitly define API endpoints
-    process_translation_api = gr.Interface(
-        fn=process_translation,
-        inputs=[gr.Textbox(label="text"), gr.Radio(label="model"), gr.Dropdown(label="src_lang"), gr.Dropdown(label="tgt_lang")],
-        outputs=gr.Textbox(label="translation_output"),
-        api_name="process_translation"
-    )
-
-    process_tts_api = gr.Interface(
-        fn=process_tts,
-        inputs=gr.Textbox(label="text"),
-        outputs=gr.Audio(label="audio_output"),
-        api_name="process_tts"
-    )
-
-    process_transcription_api = gr.Interface(
-        fn=process_transcription,
-        inputs=[gr.Audio(label="audio_path"), gr.Radio(label="language")],
-        outputs=gr.Textbox(label="transcription_output"),
-        api_name="process_transcription"
-    )
-
-    # Add API endpoints to the app
-    process_translation_api.render()
-    process_tts_api.render()
-    process_transcription_api.render()
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()