update llama custom

models/llamaCustom.py

@@ -20,7 +20,6 @@ from llama_index import (
     StorageContext,
     load_index_from_storage,
 )
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from llama_index.llms import CompletionResponse, CustomLLM, LLMMetadata
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
@@ -58,27 +57,6 @@ def load_model(model_name: str):
     return pipe
 
 
-@st.cache_resource
-def load_model(mode_name: str):
-    # llm_model_name = "bigscience/bloom-560m"
-    tokenizer = AutoTokenizer.from_pretrained(mode_name)
-    model = AutoModelForCausalLM.from_pretrained(mode_name, config="T5Config")
-
-    pipe = pipeline(
-        task="text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        # device=0, # GPU device number
-        # max_length=512,
-        do_sample=True,
-        top_p=0.95,
-        top_k=50,
-        temperature=0.7,
-    )
-
-    return pipe
-
-
 class OurLLM(CustomLLM):
     def __init__(self, model_name: str, model_pipeline):
         self.model_name = model_name
@@ -104,10 +82,12 @@ class OurLLM(CustomLLM):
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         raise NotImplementedError()
 
+
 class LlamaCustom:
     def __init__(self, model_name: str) -> None:
         self.vector_index = self.initialize_index(model_name=model_name)
 
+    @st.cache_resource
     def initialize_index(_self, model_name: str):
         index_name = model_name.split("/")[-1]