Update src/streamlit_app.py

src/streamlit_app.py

@@ -4,7 +4,8 @@ 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"
+# 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"
 
 st.set_page_config(
     page_icon="💬",
@@ -14,8 +15,13 @@ 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
 
 
@@ -26,7 +32,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()