update llama custom

models/llamaCustom.py

@@ -58,34 +58,79 @@ def load_model(model_name: str):
     return pipe
 
 
-class CustomLLM(LLM):
-    llm_model_name: str
-    pipeline: Any
+@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")
 
-    def __init__(self, llm_model_name: str):
-        super().__init__()
-        
-        self.llm_model_name = llm_model_name
-        self.pipeline = load_model(mode_name=llm_model_name)
+    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,
+    )
 
-    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+    return pipe
+
+
+class OurLLM(CustomLLM):
+    def __init__(self, model_name: str, model_pipeline):
+        self.model_name = model_name
+        self.pipeline = model_pipeline
+
+    @property
+    def metadata(self) -> LLMMetadata:
+        """Get LLM metadata."""
+        return LLMMetadata(
+            context_window=CONTEXT_WINDOW,
+            num_output=NUM_OUTPUT,
+            model_name=self.model_name,
+        )
+
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         prompt_length = len(prompt)
-        response = self.pipeline(prompt, max_new_tokens=525)[0]["generated_text"]
+        response = self.pipeline(prompt, max_new_tokens=NUM_OUTPUT)[0]["generated_text"]
 
         # only return newly generated tokens
-        return response[prompt_length:]
+        text = response[prompt_length:]
+        return CompletionResponse(text=text)
 
-    @property
-    def _identifying_params(self) -> Mapping[str, Any]:
-        return {"name_of_model": self.llm_model_name}
+    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        raise NotImplementedError()
 
-    @property
-    def _llm_type(self) -> str:
-        return "custom"
+    # 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
@@ -93,6 +138,7 @@ class LlamaCustom:
         index_name = model_name.split("/")[-1]
 
         file_path = f"./vectorStores/{index_name}"
+
         if os.path.exists(path=file_path):
             # rebuild storage context
             storage_context = StorageContext.from_defaults(persist_dir=file_path)
@@ -105,23 +151,11 @@ class LlamaCustom:
             #     index = pickle.loads(file.readlines())
             return index
         else:
-            # define llm
-            prompt_helper = PromptHelper(
-                context_window=CONTEXT_WINDOW,
-                num_output=NUM_OUTPUT,
-                chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
-            )
-
-            llm_predictor = LLMPredictor(llm=CustomLLM(llm_model_name=model_name))
-            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=service_context
+                documents, service_context=self.service_context
             )
 
             # local write access
@@ -134,6 +168,10 @@ 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()
+        query_engine = self.vector_index.as_query_engine(
+            text_qa_template=text_qa_template, refine_template=refine_template
+        )
         response = query_engine.query(query_str)
+        print("metadata: ", response.metadata)
         return str(response)