added streamlit cache

app.py

@@ -20,26 +20,37 @@ from qdrant_client import QdrantClient
 from config import DB_CONFIG, DB_E5_CONFIG
 
 
-E5_MODEL_NAME = "intfloat/multilingual-e5-large"
-E5_MODEL_KWARGS = {"device": "cuda:0" if torch.cuda.is_available() else "cpu"}
-E5_ENCODE_KWARGS = {"normalize_embeddings": False}
-E5_EMBEDDINGS = HuggingFaceEmbeddings(
-    model_name=E5_MODEL_NAME,
-    model_kwargs=E5_MODEL_KWARGS,
-    encode_kwargs=E5_ENCODE_KWARGS,
-)
-
-if False and torch.cuda.is_available():  # TODO: for local debug
-    RINNA_MODEL_NAME = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
-    RINNA_TOKENIZER = AutoTokenizer.from_pretrained(RINNA_MODEL_NAME, use_fast=False)
-    RINNA_MODEL = AutoModelForCausalLM.from_pretrained(
-        RINNA_MODEL_NAME,
-        load_in_8bit=True,
-        torch_dtype=torch.float16,
-        device_map="auto",
+@st.cache_resource
+def load_e5_embeddings():
+    model_name = "intfloat/multilingual-e5-large"
+    model_kwargs = {"device": "cuda:0" if torch.cuda.is_available() else "cpu"}
+    encode_kwargs = {"normalize_embeddings": False}
+    embeddings = HuggingFaceEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs,
     )
-else:
-    RINNA_MODEL = None
+    return embeddings
+
+
+@st.cache_resource
+def load_rinna_model():
+    if torch.cuda.is_available():
+        model_name = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            load_in_8bit=True,
+            torch_dtype=torch.float16,
+            device_map="auto",
+        )
+        return tokenizer, model
+    else:
+        return None, None
+
+
+E5_EMBEDDINGS = load_e5_embeddings()
+RINNA_TOKENIZER, RINNA_MODEL = load_rinna_model()
 
 
 def _get_config_and_embeddings(collection_name: str | None) -> tuple:
@@ -54,7 +65,8 @@ def _get_config_and_embeddings(collection_name: str | None) -> tuple:
     return db_config, embeddings
 
 
-def _get_rinna_llm(temperature: float):
+@st.cache_resource
+def _get_rinna_llm(temperature: float) -> HuggingFacePipeline | None:
     if RINNA_MODEL is not None:
         pipe = pipeline(
             "text-generation",
@@ -95,7 +107,7 @@ def get_retrieval_qa(
     model_name: str | None,
     temperature: float,
     option: str | None,
-) -> RetrievalQA:
+):
     db_config, embeddings = _get_config_and_embeddings(collection_name)
     db_url, db_api_key, db_collection_name = db_config
     client = QdrantClient(url=db_url, api_key=db_api_key)
@@ -125,7 +137,7 @@ def get_retrieval_qa(
     return result
 
 
-def get_related_url(metadata):
+def get_related_url(metadata) -> Iterable[str]:
     urls = set()
     for m in metadata:
         # p = m['source']

requirements.txt

@@ -10,3 +10,4 @@ accelerate
 bitsandbytes
 scipy
 sentence_transformers
+streamlit