Update src/streamlit_app.py

src/streamlit_app.py

@@ -4,8 +4,7 @@ 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"
-torch.backends.quantized.engine = "qnnpack" if torch.mps.is_available() else "fbgemm"
+DEVICE = "mps" if torch.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
 
 st.set_page_config(
     page_icon="💬",
@@ -15,13 +14,8 @@ st.set_page_config(
 
 @st.cache_resource
 def get_model() -> tuple[AutoModelForCausalLM, AutoTokenizer]:
-    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)#.to(DEVICE)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH).to(DEVICE)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-
-    model = torch.quantization.quantize_dynamic(
-        model, {torch.nn.Linear}, dtype=torch.qint8
-    )
-
     return model, tokenizer
 
 
@@ -32,7 +26,7 @@ with st.spinner(text="Loading model, please wait...", show_time=True):
 @torch.inference_mode()
 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)
+    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).strip()
@@ -57,4 +51,4 @@ else:
 
 if query := st.text_input("Query", placeholder=f"Write in {src}"):
     with st.spinner():
-        st.text(translate(src, tgt, query))
+        st.text(translate(src, tgt, query))