Update app.py

app.py

@@ -1,90 +1,90 @@
-import streamlit as st
-from transformers import (
-    PreTrainedTokenizerBase,
-    PreTrainedTokenizerFast,
-    AutoModelForCausalLM,
-)
-
-model_dict = {
-    "NanoTranslator-XS": "Mxode/NanoTranslator-XS",
-    "NanoTranslator-S": "Mxode/NanoTranslator-S",
-    "NanoTranslator-M": "Mxode/NanoTranslator-M",
-    "NanoTranslator-M2": "Mxode/NanoTranslator-M2",
-    "NanoTranslator-L": "Mxode/NanoTranslator-L",
-    "NanoTranslator-XL": "Mxode/NanoTranslator-XL",
-    "NanoTranslator-XXL": "Mxode/NanoTranslator-XXL",
-    "NanoTranslator-XXL2": "Mxode/NanoTranslator-XXL2",
-}
-
-
-# initialize model
-@st.cache_resource
-def load_model(model_path: str):
-    model = AutoModelForCausalLM.from_pretrained(model_path)
-    tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
-    return model, tokenizer
-
-
-def translate(text: str, model, tokenizer: PreTrainedTokenizerBase, **kwargs):
-    generation_args = dict(
-        max_new_tokens=kwargs.pop("max_new_tokens", 64),
-        do_sample=kwargs.pop("do_sample", True),
-        temperature=kwargs.pop("temperature", 0.55),
-        top_p=kwargs.pop("top_p", 0.8),
-        top_k=kwargs.pop("top_k", 40),
-        **kwargs
-    )
-
-    prompt = "<|im_start|>" + text + "<|endoftext|>"
-    model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
-
-    generated_ids = model.generate(model_inputs.input_ids, **generation_args)
-    generated_ids = [
-        output_ids[len(input_ids) :]
-        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
-
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return response
-
-
-st.title("NanoTranslator-Demo")
-
-st.sidebar.title("Options")
-model_choice = st.sidebar.selectbox("Model", list(model_dict.keys()))
-do_sample = st.sidebar.checkbox("do_sample", value=True)
-max_new_tokens = st.sidebar.slider(
-    "max_new_tokens", min_value=1, max_value=256, value=64
-)
-temperature = st.sidebar.slider(
-    "temperature", min_value=0.01, max_value=1.5, value=0.55, step=0.01
-)
-top_p = st.sidebar.slider("top_p", min_value=0.01, max_value=1.0, value=0.8, step=0.01)
-top_k = st.sidebar.slider("top_k", min_value=1, max_value=100, value=40, step=1)
-
-# 根据选择的模型加载
-model_path = model_dict[model_choice]
-model, tokenizer = load_model(model_path)
-
-input_text = st.text_area(
-    "Please input the text to be translated (Currently supports only English to Chinese):",
-    "Each step of the cell cycle is monitored by internal.",
-)
-
-if st.button("translate"):
-    if input_text.strip():
-        with st.spinner("Translating..."):
-            translation = translate(
-                input_text,
-                model,
-                tokenizer,
-                max_new_tokens=max_new_tokens,
-                do_sample=do_sample,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-            )
-        st.success("Translated successfully!")
-        st.write(translation)
-    else:
-        st.warning("Please input text before translation!")
+import streamlit as st
+from transformers import (
+    PreTrainedTokenizerBase,
+    PreTrainedTokenizerFast,
+    AutoModelForCausalLM,
+)
+
+model_dict = {
+    "NanoTranslator-XS": "Mxode/NanoTranslator-XS",
+    "NanoTranslator-S": "Mxode/NanoTranslator-S",
+    "NanoTranslator-M": "Mxode/NanoTranslator-M",
+    "NanoTranslator-M2": "Mxode/NanoTranslator-M2",
+    "NanoTranslator-L": "Mxode/NanoTranslator-L",
+    "NanoTranslator-XL": "Mxode/NanoTranslator-XL",
+    "NanoTranslator-XXL": "Mxode/NanoTranslator-XXL",
+    "NanoTranslator-XXL2": "Mxode/NanoTranslator-XXL2",
+}
+
+
+# initialize model
+@st.cache_resource
+def load_model(model_path: str):
+    model = AutoModelForCausalLM.from_pretrained(model_path)
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
+    return model, tokenizer
+
+
+def translate(text: str, model, tokenizer: PreTrainedTokenizerBase, **kwargs):
+    generation_args = dict(
+        max_new_tokens=kwargs.pop("max_new_tokens", 64),
+        do_sample=kwargs.pop("do_sample", True),
+        temperature=kwargs.pop("temperature", 0.55),
+        top_p=kwargs.pop("top_p", 0.8),
+        top_k=kwargs.pop("top_k", 40),
+        **kwargs
+    )
+
+    prompt = "<|im_start|>" + text + "<|endoftext|>"
+    model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+
+    generated_ids = model.generate(model_inputs.input_ids, **generation_args)
+    generated_ids = [
+        output_ids[len(input_ids) :]
+        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return response
+
+
+st.title("NanoTranslator-Demo")
+
+st.sidebar.title("Options")
+model_choice = st.sidebar.selectbox("Model", list(model_dict.keys()), index=list(model_options.keys()).index("NanoTranslator-XXL2"))
+do_sample = st.sidebar.checkbox("do_sample", value=True)
+max_new_tokens = st.sidebar.slider(
+    "max_new_tokens", min_value=1, max_value=256, value=64
+)
+temperature = st.sidebar.slider(
+    "temperature", min_value=0.01, max_value=1.5, value=0.55, step=0.01
+)
+top_p = st.sidebar.slider("top_p", min_value=0.01, max_value=1.0, value=0.8, step=0.01)
+top_k = st.sidebar.slider("top_k", min_value=1, max_value=100, value=40, step=1)
+
+# 根据选择的模型加载
+model_path = model_dict[model_choice]
+model, tokenizer = load_model(model_path)
+
+input_text = st.text_area(
+    "Please input the text to be translated (Currently supports only English to Chinese):",
+    "Each step of the cell cycle is monitored by internal.",
+)
+
+if st.button("translate"):
+    if input_text.strip():
+        with st.spinner("Translating..."):
+            translation = translate(
+                input_text,
+                model,
+                tokenizer,
+                max_new_tokens=max_new_tokens,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+            )
+        st.success("Translated successfully!")
+        st.write(translation)
+    else:
+        st.warning("Please input text before translation!")