v1.2.1 translate model done

app.py

@@ -1,51 +1,60 @@
 import streamlit as st
-
 from transformers import T5ForConditionalGeneration, T5Tokenizer
+from transformers import pipeline
+import torch
+import scipy
+
+st.title("FinalProject")
+
+
+@st.cache_resource
+def load_summarization_model():
+    print("Loading summarization model...")
+    return pipeline("summarization", model="facebook/bart-large-cnn")
+
+summarizer = load_summarization_model()
+
+ARTICLE = st.text_area("Enter the article to summarize:", height=300)
 
-device = 'cpu' #or 'cpu' for translate on cpu
+max_length = st.number_input("Enter max length for summary:", min_value=10, max_value=500, value=130)
+min_length = st.number_input("Enter min length for summary:", min_value=5, max_value=450, value=30)
 
-model_name = 'utrobinmv/t5_translate_en_ru_zh_large_1024'
-model = T5ForConditionalGeneration.from_pretrained(model_name)
-model.to(device)
-tokenizer = T5Tokenizer.from_pretrained(model_name)
 
-prefix = 'translate to en: '
-src_text = prefix + "Съешь ещё этих мягких французских булок."
+device = 'cpu'
 
-# translate Russian to Chinese
-input_ids = tokenizer(src_text, return_tensors="pt")
+@st.cache_resource
+def load_translation_model():
+    model_name = 'utrobinmv/t5_translate_en_ru_zh_large_1024'
+    model = T5ForConditionalGeneration.from_pretrained(model_name)
+    model.to(device)
+    return model, T5Tokenizer.from_pretrained(model_name)
 
-generated_tokens = model.generate(**input_ids.to(device))
 
-result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
-print(result)
-st.write(result[0])
-# 再吃这些法国的甜蜜的面包。
-# import streamlit as st
-# from transformers import pipeline
-# import torch
-# import scipy
 
-# st.title("FinalProject")
+model, tokenizer = load_translation_model()
 
 
-# @st.cache_resource
-# def load_summarization_model():
-#     print("Loading summarization model...")
-#     return pipeline("summarization", model="facebook/bart-large-cnn")
+if st.button("Summarize"):
+    if ARTICLE.strip():
+        answer = summarizer(ARTICLE, max_length=int(max_length), min_length=int(min_length), do_sample=False)
+        summary = answer[0]['summary_text']
+        st.write("### Summary:")
+        st.write(summary)
+    else:
+        st.error("Please enter an article to summarize.")
 
-# summarizer = load_summarization_model()
+target_language = st.selectbox("Choose target language for translation:", ["ru", "zh"])
 
-# ARTICLE = st.text_area("Enter the article to summarize:", height=300)
+if st.button("Translate"):
+    if ARTICLE.strip():
+        prefix = f"translate to {target_language}: "
+        src_text = prefix + ARTICLE
 
-# max_length = st.number_input("Enter max length for summary:", min_value=10, max_value=500, value=130)
-# min_length = st.number_input("Enter min length for summary:", min_value=5, max_value=450, value=30)
+        input_ids = tokenizer(src_text, return_tensors="pt")
+        generated_tokens = model.generate(**input_ids.to(device))
+        result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
 
-# if st.button("Summarize"):
-#     if ARTICLE.strip():
-#         answer = summarizer(ARTICLE, max_length=int(max_length), min_length=int(min_length), do_sample=False)
-#         summary = answer[0]['summary_text']
-#         st.write("### Summary:")
-#         st.write(summary)
-#     else:
-#         st.error("Please enter an article to summarize.")
+        st.write(f"### Translation ({target_language.upper()}):")
+        st.write(result[0])
+    else:
+        st.error("Please enter an article to translate.")