Update models/llamaCustom.py

models/llamaCustom.py

@@ -6,9 +6,9 @@ from typing import Any, List, Mapping, Optional
 import numpy as np
 import openai
 import pandas as pd
-import streamlit as st
 from dotenv import load_dotenv
-from huggingface_hub import HfFileSystem, Repository
+from huggingface_hub import HfFileSystem
+from langchain.llms.base import LLM
 from llama_index import (
     Document,
     GPTVectorStoreIndex,
@@ -19,17 +19,12 @@ from llama_index import (
     StorageContext,
     load_index_from_storage,
 )
-from llama_index.llms import CompletionResponse, CustomLLM, LLMMetadata
-
-# from langchain.llms.base import LLM
-# from llama_index.prompts import Prompt
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # from utils.customLLM import CustomLLM
 
 load_dotenv()
 # openai.api_key = os.getenv("OPENAI_API_KEY")
-
 fs = HfFileSystem()
 
 # define prompt helper
@@ -38,98 +33,62 @@ CONTEXT_WINDOW = 2048
 # set number of output tokens
 NUM_OUTPUT = 525
 # set maximum chunk overlap
-CHUNK_OVERLAP_RATIO = 0.2
+CHUNK_OVERLAP_RATION = 0.2
 
 prompt_helper = PromptHelper(
     context_window=CONTEXT_WINDOW,
     num_output=NUM_OUTPUT,
-    chunk_overlap_ratio=CHUNK_OVERLAP_RATIO,
+    chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
 )
 
-
-@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
+llm_model_name = "bigscience/bloom-560m"
+tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
+model = AutoModelForCausalLM.from_pretrained(llm_model_name, config="T5Config")
+
+model_pipeline = pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    task="text-generation",
+    # device=0, # GPU device number
+    # max_length=512,
+    do_sample=True,
+    top_p=0.95,
+    top_k=50,
+    temperature=0.7,
+)
 
 
-class OurLLM(CustomLLM):
-    def __init__(self, model_name: str, model_pipeline):
-        self.model_name = model_name
-        self.pipeline = model_pipeline
+class CustomLLM(LLM):
+    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:
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         prompt_length = len(prompt)
-        response = self.pipeline(prompt, max_new_tokens=NUM_OUTPUT)[0]["generated_text"]
+        response = self.pipeline(prompt, max_new_tokens=525)[0]["generated_text"]
 
         # only return newly generated tokens
-        text = response[prompt_length:]
-        return CompletionResponse(text=text)
-
-    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        raise NotImplementedError()
+        return response[prompt_length:]
 
-    # 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 _identifying_params(self) -> Mapping[str, Any]:
+        return {"name_of_model": self.model_name}
 
-    # @property
-    # def _llm_type(self) -> str:
-    #     return "custom"
+    @property
+    def _llm_type(self) -> str:
+        return "custom"
 
 
-@st.cache_resource
 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)
-
-    def initialize_index(self, model_name: str):
-        index_name = model_name.split("/")[-1]
+    llm_predictor = LLMPredictor(llm=CustomLLM())
+    service_context = ServiceContext.from_defaults(
+        llm_predictor=llm_predictor, prompt_helper=prompt_helper
+    )
 
-        file_path = f"./vectorStores/{index_name}"
+    def __init__(self, name: str) -> None:
+        self.vector_index = self.initialize_index(index_name=name)
 
+    def initialize_index(self, index_name):
+        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)
@@ -160,9 +119,5 @@ 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
-        # )
         response = query_engine.query(query_str)
-        print("metadata: ", response.metadata)
-        return str(response)
+        return str(response)