Update app.py

app.py

@@ -1,69 +1,56 @@
 import streamlit as st
-import torch
-from transformers import AutoModelForSeq2SeqLM, NllbTokenizer, T5Tokenizer, T5ForConditionalGeneration
-from transformers import pipeline
+from transformers import AutoModelForSeq2SeqLM, T5ForConditionalGeneration
+from transformers import NllbTokenizer, T5Tokenizer
 
-# Load translation model and tokenizer (same as before)
-model_load_name = 'sarahai/nllb-uzbek-cyrillic-to-russian'
-model = AutoModelForSeq2SeqLM.from_pretrained(model_load_name)
-tokenizer = NllbTokenizer.from_pretrained(model_load_name)
+# Load translation model and tokenizer
+translation_model_name = 'sarahai/nllb-uzbek-cyrillic-to-russian'
+translation_model = AutoModelForSeq2SeqLM.from_pretrained(translation_model_name)
+translation_tokenizer = NllbTokenizer.from_pretrained(translation_model_name)
 
+# Load summarization model and tokenizer
+summarization_model_name = 'sarahai/ruT5-base-summarizer'
+summarization_model = T5ForConditionalGeneration.from_pretrained(summarization_model_name)
+summarization_tokenizer = T5Tokenizer.from_pretrained(summarization_model_name)
 
-# Define translation function (same as before)
-def translate(text, model, tokenizer, src_lang='uzb_Cyrl', tgt_lang='rus_Cyrl', a=16, b=1.5, max_input_length=1024):
+def translate(text, model, tokenizer, src_lang='uzb_Cyrl', tgt_lang='rus_Cyrl', max_input_length=1024):
     tokenizer.src_lang = src_lang
     tokenizer.tgt_lang = tgt_lang
     inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=max_input_length)
-    result = model.generate(
-        **inputs,
-        forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
-        max_new_tokens=int(a + b * inputs.input_ids.shape[1])
-        
+    outputs = model.generate(
+        inputs['input_ids'],
+        forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang],
+        max_length=512
     )
-    return tokenizer.batch_decode(result, skip_special_tokens=True)
-
-# Load summarization model and tokenizer
-model_name = "sarahai/ruT5-base-summarizer"
-summarizer_tokenizer = T5Tokenizer.from_pretrained(model_name)
-summarizer_model = T5ForConditionalGeneration.from_pretrained(model_name)
-
-# Define summarization function (using model and tokenizer)
-def summarize(russian_text):
-    try:
-        input_ids = summarizer_tokenizer(text, return_tensors="pt", padding="max_length").input_ids
-        outputs = summarizer_model.generate(
-            input_ids,
-            max_length=250,
-            min_length=150,
-            length_penalty=2.0,
-            num_beams=4,
-            early_stopping=True,
-        )
-        summary = summarizer_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return summary
-    except Exception as e:
-        return f"Ошибка: {str(e)}"
-
-# Streamlit interface
-st.title("Перевод и аннотация текста")
-text = st.text_area("Введите текст на узбекском языке", height=200)
-
-if st.button("Перевести и аннотировать"):
-    if text:
-        try:
-            # Translate Uzbek text to Russian
-            russian_text = translate(text, model, tokenizer, src_lang='uzb_Cyrl', tgt_lang='rus_Cyrl', a=16, b=1.5, max_input_length=1024)
-
-            # Summarize the translated Russian text
-            summary = summarize(russian_text)
-
-            # Display results
-            st.success("Перевод на русский:")
-            st.write(russian_text)
-            st.success("Аннотация русского текста:")
-            st.write(summary)
-        except Exception as e:
-            st.error(f"Ошибка: {e}")
+    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return translated_text
+
+def summarize(translated_text, model, tokenizer, max_length=250, min_length=150):
+    input_ids = tokenizer.encode("summarize: " + translated_text, return_tensors="pt", max_length=1024, truncation=True)
+    summary_ids = model.generate(
+        input_ids,
+        max_length=max_length,
+        min_length=min_length,
+        length_penalty=2.0,
+        num_beams=4,
+        early_stopping=True
+    )
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
+
+# Streamlit app setup
+st.title("Russian to Uzbek Translation and Summarization")
+input_text = st.text_area("Enter text in Russian:", height=200)
+
+if st.button("Translate and Summarize"):
+    if input_text:
+        with st.spinner('Translating...'):
+            translated_text = translate(input_text, translation_model, translation_tokenizer)
+            st.text_area("Translated Text (Uzbek):", value=translated_text, height=200)
+        
+        with st.spinner('Summarizing...'):
+            summary_text = summarize(translated_text, summarization_model, summarization_tokenizer, max_length=250, min_length=150)
+            st.text_area("Summary (Uzbek):", value=summary_text, height=100)
     else:
-        st.warning("Пожалуйста, введите текст на узбекском языке.")
+        st.warning("Please enter text in Russian to proceed.")
+