Update models/llamaCustom.py

models/llamaCustom.py

@@ -5,6 +5,7 @@ from typing import Any, List, Mapping, Optional
 
 import numpy as np
 import openai
+import streamlit as st
 import pandas as pd
 from dotenv import load_dotenv
 from huggingface_hub import HfFileSystem
@@ -76,7 +77,7 @@ class CustomLLM(LLM):
     def _llm_type(self) -> str:
         return "custom"
 
-
+@st.cache_resource
 class LlamaCustom:
     # define llm
     llm_predictor = LLMPredictor(llm=CustomLLM())
@@ -84,10 +85,12 @@ class LlamaCustom:
         llm_predictor=llm_predictor, prompt_helper=prompt_helper
     )
 
-    def __init__(self, name: str) -> None:
-        self.vector_index = self.initialize_index(index_name=name)
+    def __init__(self, model_name: str) -> None:
+        self.vector_index = self.initialize_index(index_name=model_name)
 
-    def initialize_index(self, index_name):
+    def initialize_index(self, model_name: str):
+        index_name = model_name.split("/")[-1]
+        
         file_path = f"./vectorStores/{index_name}"
         if os.path.exists(path=file_path):
             # rebuild storage context