Update models/llamaCustom.py

models/llamaCustom.py

@@ -36,12 +36,6 @@ NUM_OUTPUT = 525
 # set maximum chunk overlap
 CHUNK_OVERLAP_RATION = 0.2
 
-prompt_helper = PromptHelper(
-    context_window=CONTEXT_WINDOW,
-    num_output=NUM_OUTPUT,
-    chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
-)
-
 llm_model_name = "bigscience/bloom-560m"
 tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
 model = AutoModelForCausalLM.from_pretrained(llm_model_name, config="T5Config")
@@ -71,7 +65,7 @@ class CustomLLM(LLM):
 
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
-        return {"name_of_model": self.model_name}
+        return {"name_of_model": llm_model_name}
 
     @property
     def _llm_type(self) -> str:
@@ -80,6 +74,11 @@ class CustomLLM(LLM):
 @st.cache_resource
 class LlamaCustom:
     # define llm
+    prompt_helper = PromptHelper(
+        context_window=CONTEXT_WINDOW,
+        num_output=NUM_OUTPUT,
+        chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
+    )
     llm_predictor = LLMPredictor(llm=CustomLLM())
     service_context = ServiceContext.from_defaults(
         llm_predictor=llm_predictor, prompt_helper=prompt_helper