update llama custom

models/llamaCustom.py

@@ -21,8 +21,8 @@ from llama_index import (
     load_index_from_storage,
 )
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-
-# from utils.customLLM import CustomLLM
+from llama_index.llms import CompletionResponse, CustomLLM, LLMMetadata
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 load_dotenv()
 # openai.api_key = os.getenv("OPENAI_API_KEY")
@@ -104,36 +104,10 @@ class OurLLM(CustomLLM):
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         raise NotImplementedError()
 
-    # def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
-    #     prompt_length = len(prompt)
-    #     response = self.pipeline(prompt, max_new_tokens=525)[0]["generated_text"]
-
-    #     # only return newly generated tokens
-    #     return response[prompt_length:]
-
-    # @property
-    # def _identifying_params(self) -> Mapping[str, Any]:
-    #     return {"name_of_model": self.model_name}
-
-    # @property
-    # def _llm_type(self) -> str:
-    #     return "custom"
-
 class LlamaCustom:
-    # define llm
-    # llm_predictor = LLMPredictor(llm=OurLLM())
-    # service_context = ServiceContext.from_defaults(
-    #     llm_predictor=llm_predictor, prompt_helper=prompt_helper
-    # )
     def __init__(self, model_name: str) -> None:
-        pipe = load_model(mode_name=model_name)
-        llm = OurLLM(model_name=model_name, model_pipeline=pipe)
-        self.service_context = ServiceContext.from_defaults(
-            llm=llm, prompt_helper=prompt_helper
-        )
         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]
 
@@ -151,11 +125,26 @@ class LlamaCustom:
             #     index = pickle.loads(file.readlines())
             return index
         else:
+            prompt_helper = PromptHelper(
+                context_window=CONTEXT_WINDOW,
+                num_output=NUM_OUTPUT,
+                chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
+            )
+
+            # define llm
+            pipe = load_model(mode_name=model_name)
+            llm = OurLLM(model_name=model_name, model_pipeline=pipe)
+
+            llm_predictor = LLMPredictor(llm=llm)
+            service_context = ServiceContext.from_defaults(
+                llm_predictor=llm_predictor, prompt_helper=prompt_helper
+            )
+
             # documents = prepare_data(r"./assets/regItems.json")
             documents = SimpleDirectoryReader(input_dir="./assets/pdf").load_data()
 
             index = GPTVectorStoreIndex.from_documents(
-                documents, service_context=self.service_context
+                documents, service_context=service_context
             )
 
             # local write access
@@ -168,10 +157,7 @@ class LlamaCustom:
 
     def get_response(self, query_str):
         print("query_str: ", query_str)
-        # query_engine = self.vector_index.as_query_engine()
-        query_engine = self.vector_index.as_query_engine(
-            text_qa_template=text_qa_template, refine_template=refine_template
-        )
+        query_engine = self.vector_index.as_query_engine()
         response = query_engine.query(query_str)
         print("metadata: ", response.metadata)
         return str(response)