Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,12 +3,17 @@ import streamlit as st
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_PATH = "NetherQuartz/tatoeba-tok-multi-gemma-2-2b-merged"
-DEVICE = "mps" if torch.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_PATH = "NetherQuartz/tatoeba-tok-multi-gemma-2-2b-merged-int4"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+st.set_page_config(
+    page_icon="💬",
+    page_title="ilo toki"
+)
 
 
 @st.cache_resource
-def get_model():
+def get_model() -> tuple[AutoModelForCausalLM, AutoTokenizer]:
     model = AutoModelForCausalLM.from_pretrained(MODEL_PATH).to(DEVICE)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
     return model, tokenizer
@@ -17,17 +22,13 @@ def get_model():
 model, tokenizer = get_model()
 
 
-def translate(src_lang, tgt_lang, query):
+def translate(src_lang: str, tgt_lang: str, query: str) -> str:
     text = f"Translate {src_lang} to {tgt_lang}.\nQuery: {query}\nAnswer:"
     tokens = tokenizer(text, return_tensors="pt").to(DEVICE)
     outputs = model.generate(**tokens)
     ans = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return ans.removeprefix(text)
+    return ans.removeprefix(text).strip()
 
-st.set_page_config(
-    page_icon="💬",
-    page_title="ilo toki"
-)
 
 st.title("💬 ilo toki")