Update app.py

app.py

@@ -59,65 +59,57 @@ LANGUAGES = {
 
 # Function to get the appropriate translation model and tokenizer
 def get_translation_model(source_lang, target_lang):
-    # Use a generic model for translation if specific model is not available
     model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
-    try:
-        model = MarianMTModel.from_pretrained(model_name)
-        tokenizer = MarianTokenizer.from_pretrained(model_name)
-    except OSError:
-        st.error(f"Model '{model_name}' not found. Please check the model name or use another language pair.")
-        return None, None
+    model = MarianMTModel.from_pretrained(model_name)
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
     return model, tokenizer
 
 # Function to translate text
 def translate_text(text, source_lang, target_lang):
     model, tokenizer = get_translation_model(source_lang, target_lang)
-    if model is None or tokenizer is None:
-        return ""
     inputs = tokenizer([text], return_tensors="pt", truncation=True)
     translated_ids = model.generate(inputs['input_ids'], max_length=1024)
     translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
     return translated_text
 
 # Summarization function with multi-language support
-def summarize_text(text, input_language="English", output_language="English"):
-    input_lang_code = LANGUAGES[input_language]
-    output_lang_code = LANGUAGES[output_language]
+def summarize_text(text, source_language="English", target_language="English"):
+    source_lang_code = LANGUAGES[source_language]
+    target_lang_code = LANGUAGES[target_language]
     
     # If the input language is not English, translate to English
-    if input_lang_code != "en_XX":
-        text = translate_text(text, input_lang_code, "en_XX")
+    if source_lang_code != "en_XX":
+        text = translate_text(text, source_lang_code, "en_XX")
     
     # Summarize the text using mBART
     inputs = multilingual_summarization_tokenizer(text, return_tensors='pt', padding=True, truncation=True)
     summary_ids = multilingual_summarization_model.generate(
         inputs['input_ids'],
-        num_beams=6,        # Increase the number of beams for better quality
-        max_length=1024,   # Increase the maximum length
-        min_length=256,    # Set a minimum length for the summary
-        length_penalty=2.0,  # Adjust length penalty to control the length of the summary
+        num_beams=6,        # Increased beams for better quality
+        max_length=1500,   # Increased maximum length for longer summaries
+        min_length=400,    # Set a minimum length for the summary
+        length_penalty=1.5,  # Adjust length penalty to control the length of the summary
         early_stopping=True
     )
     summary = multilingual_summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     
-    # Translate summary to the output language if needed
-    if output_lang_code != "en_XX":
-        summary = translate_text(summary, "en_XX", output_lang_code)
+    # Translate summary to the target language if needed
+    if target_lang_code != "en_XX":
+        summary = translate_text(summary, "en_XX", target_lang_code)
     
     return summary
 
 # Streamlit interface
 st.title("Multi-Language Text Summarization Tool")
-st.write("Enter the text you want to summarize, select the input language, and choose the output language for the summary.")
 
-text_input = st.text_area("Input Text")
-input_language = st.selectbox("Input Language", options=list(LANGUAGES.keys()), index=list(LANGUAGES.keys()).index("English"))
-output_language = st.selectbox("Output Language", options=list(LANGUAGES.keys()), index=list(LANGUAGES.keys()).index("English"))
+text = st.text_area("Input Text")
+source_language = st.selectbox("Source Language", options=list(LANGUAGES.keys()), index=list(LANGUAGES.keys()).index("English"))
+target_language = st.selectbox("Target Language", options=list(LANGUAGES.keys()), index=list(LANGUAGES.keys()).index("English"))
 
 if st.button("Summarize"):
-    if text_input:
-        summary = summarize_text(text_input, input_language, output_language)
-        st.write("Summary:")
+    if text:
+        summary = summarize_text(text, source_language, target_language)
+        st.subheader("Summary")
         st.write(summary)
     else:
-        st.warning("Please enter some text to summarize.")
+        st.warning("Please enter text to summarize.")