updated code for llamaCustom

models/langOpen.py

@@ -4,10 +4,14 @@ import openai
 from dotenv import load_dotenv
 from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
+from langchain.document_loaders import PyPDFLoader
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain.vectorstores import FAISS
 
+loader = PyPDFLoader("./assets/pdf/CADWReg.pdf")
+pages = loader.load_and_split()
+
 load_dotenv()
 
 embeddings = OpenAIEmbeddings()
@@ -31,9 +35,7 @@ class LangOpen:
         if os.path.exists(path=path):
             return FAISS.load_local(folder_path=path, embeddings=embeddings)
         else:
-            faiss = FAISS.from_texts(
-                "./assets/updated_calregs.txt", embedding=embeddings
-            )
+            faiss = FAISS.from_documents(pages, embeddings)
             faiss.save_local(path)
             return faiss
 

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
-from langchain.llms.base import LLM
+from huggingface_hub import HfFileSystem, Repository
 from llama_index import (
     Document,
     GPTVectorStoreIndex,
@@ -19,12 +19,17 @@ from llama_index import (
     StorageContext,
     load_index_from_storage,
 )
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+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 utils.customLLM import CustomLLM
 
 load_dotenv()
 # openai.api_key = os.getenv("OPENAI_API_KEY")
+
 fs = HfFileSystem()
 
 # define prompt helper
@@ -33,62 +38,122 @@ CONTEXT_WINDOW = 2048
 # set number of output tokens
 NUM_OUTPUT = 525
 # set maximum chunk overlap
-CHUNK_OVERLAP_RATION = 0.2
+CHUNK_OVERLAP_RATIO = 0.2
 
 prompt_helper = PromptHelper(
     context_window=CONTEXT_WINDOW,
     num_output=NUM_OUTPUT,
-    chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
+    chunk_overlap_ratio=CHUNK_OVERLAP_RATIO,
 )
 
-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,
+text_qa_template_str = (
+    "Context information is below.\n"
+    "---------------------\n"
+    "{context_str}\n"
+    "---------------------\n"
+    "Using both the context information and also using your own knowledge, "
+    "answer the question: {query_str}\n"
+    "If the question is relevant, you can answer by providing the name of the chapter, the article and the title to the answer. In addition, you can add the page number of the document when you found the answer.\n"
+    "If the context isn't helpful, you can also answer the question on your own.\n"
 )
+text_qa_template = Prompt(text_qa_template_str)
+
+refine_template_str = (
+    "The original question is as follows: {query_str}\n"
+    "We have provided an existing answer: {existing_answer}\n"
+    "We have the opportunity to refine the existing answer "
+    "(only if needed) with some more context below.\n"
+    "------------\n"
+    "{context_msg}\n"
+    "------------\n"
+    "Using both the new context and your own knowledege, update or repeat the existing answer.\n"
+)
+refine_template = Prompt(refine_template_str)
+
+
+@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 CustomLLM(LLM):
-    pipeline = model_pipeline
 
-    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+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.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"
 
 
+@st.cache_resource
 class LlamaCustom:
     # define llm
-    llm_predictor = LLMPredictor(llm=CustomLLM())
-    service_context = ServiceContext.from_defaults(
-        llm_predictor=llm_predictor, prompt_helper=prompt_helper
-    )
-
-    def __init__(self, name: str) -> None:
-        self.vector_index = self.initialize_index(index_name=name)
+    # 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]
 
-    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)
@@ -118,6 +183,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)

pages/langchain_demo.py

@@ -17,7 +17,7 @@ if "openai_api_key" not in st.session_state:
     st.info("Enter your openai key to access the chatbot.")
 else:
     option = st.selectbox(
-        label="Select your model:", options=("gpt-3.5-turbo", "gpt-4"), index=0
+        label="Select your model:", options=("gpt-3.5-turbo", "gpt-4")
     )
 
     with st.spinner(f"Initializing {option} ..."):

pages/llama_custom_demo.py

@@ -1,11 +1,10 @@
 import os
-import time
 
 import openai
 import streamlit as st
 
 from models.llamaCustom import LlamaCustom
-from utils.chatbox import *
+from utils.chatbox import chatbox
 
 st.set_page_config(page_title="Llama", page_icon="🦙")
 
@@ -17,7 +16,11 @@ if "messages" not in st.session_state:
 if "openai_api_key" not in st.session_state:
     st.info("Enter your openai key to access the chatbot.")
 else:
+    option = st.selectbox(
+        label="Select your model:", options=("bigscience/bloom-560m",)
+    )
+
     with st.spinner("Initializing vector index"):
-        model = LlamaCustom(name="llamaCustom")
+        model = LlamaCustom(model_name=option)
 
     chatbox("llama_custom", model)

utils/chatbox.py

@@ -40,7 +40,6 @@ def display_bot_msg(model_name: str, bot_response: str):
         {"model_name": model_name, "role": "assistant", "content": full_response}
     )
 
-# @st.cache_data
 def chatbox(model_name: str, model: None):
     # Display chat messages from history on app rerun
     for message in st.session_state.messages: