minor changes

app.py

@@ -27,34 +27,35 @@ def detect_language(text):
     return lang_code
 
 
-def translate_to_english(text):
-    # Always translate to English (en_XX)
-    mbart_lang_code = "en_XX"
+def translate_to_english(text, src_lang):
+    # Define the target language as English
+    tgt_lang = "en_XX"
     
-    # Encode the input text for translation
+    # Tokenize the input text with the appropriate source and target language tokens
     inputs = tokenizer(
-        f"<{mbart_lang_code}>{text}", 
+        text, 
         return_tensors="pt", 
         max_length=1024, 
         truncation=True
     )
-    
-    # Perform the translation
-    translated_ids = translator.model.generate(
+
+    # Specify the source language and target language in the generation call
+    translated_ids = model.generate(
         inputs["input_ids"], 
         max_length=100, 
         length_penalty=2.0, 
-        num_beams=4
+        num_beams=4,
+        decoder_start_token_id=tokenizer.lang_code_to_id[tgt_lang],  # Explicitly set the target language
+        forced_bos_token_id=tokenizer.lang_code_to_id[src_lang]  # Set the source language
     )
     
     # Decode the translated text
     translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
-    
-    # Remove any special language code tokens like "<en_XX>"
     translated_text = re.sub(r"<[^>]+>", "", translated_text).strip()
     
     return translated_text
 
+
 def summarize_text(text, lang_code):
     mbart_lang_code = LANGUAGE_CODES.get(lang_code, "en_XX")  # Default to English if unsupported
     inputs = tokenizer(
@@ -71,9 +72,10 @@ def summarize_text(text, lang_code):
         num_beams=4
     )
     summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    summary = summary.replace(f"<>", "").strip()
+    summary = re.sub(r"<[^>]+>", "", summary).strip()
     return summary
 
+
 st.title("Multilingual Summarization and Translation App")
 st.markdown("""This app detects the language of the input text, summarizes it in the same language, and translates it into English.""")
 
@@ -95,7 +97,7 @@ if st.button("Process Text"):
                 st.write(summary)
 
                 # Then translate the summary to English
-                translation = translate_to_english(summary)
+                translation = translate_to_english(summary, LANGUAGE_CODES.get(lang_code, "en_XX"))
                 st.write("### Translated Text (English):")
                 st.write(translation)