v0.1

app.py

@@ -1,63 +1,74 @@
 import os
 import time
-import itertools
-from typing import Literal
 import gradio as gr
 import google.generativeai as genai
+from src.llamaindex_palm import LlamaIndexPaLM
+
+import logging
+logging.basicConfig(format='%(asctime)s %(clientip)-15s %(user)-8s %(message)s', level=logging.INFO)
+logger = logging.getLogger('llm')
+
+# Llama-Index LLM
+llm = LlamaIndexPaLM()
+llm.set_index_from_pinecone()
 
 # Credentials
 genai.configure(api_key=os.getenv('PALM_API_KEY'))
 
 # Gradio
-chat_defaults = {
-  'model': 'models/chat-bison-001',
-  'temperature': 0.25,
-  'candidate_count': 1,
-  'top_k': 40,
-  'top_p': 0,
-}
-
 chat_history = []
 
-def clear_chat():
+def clear_chat() -> None:
+    global chat_history 
     chat_history = []
     return None
 
-def generate_chat(prompt: str, chat_messages=chat_history):
-    print(chat_messages)
-    context = """
-    You are a proxy to Gerard Lee to answer questions as himself. 
-    I am chatting with someone who might interest in my background. I am going to reply their messages given the context below and without hallucinations.
+def generate_chat(prompt: str, llamaindex_llm: LlamaIndexPaLM):
+    global chat_history
+    # get chat history
+    context_chat_history = "\n".join(chat_history)
+
+    logger.info("Generating Message...")
+    logger.info(f"User Message:\n{prompt}\n")
+    chat_history.append(prompt)
+
+    # get context
+    context_from_index = llamaindex_llm.generate_response(prompt)
+    logger.info(f"Context from Llama-Index:\n{context_from_index}\n")
+
+    prompt_with_context = f"""
+    Rule: 
+        You are in a role play of Gerard Lee and you need to pretend to be him to answer questions from people who interested in Gerard's background.
+        You are going to reply their messages given the context below and without hallucinations. If you don't know the answer, simply say "I have no idea how to answer this question".
 
-    Context: '
-        I am a data enthusiast with more than 5 years of experience on data analytics domain. Currently working under DBS as a data scientist. 
-        I drive NLP and ML use cases also lead 2 contract data analysts to deliver analytical solutions like developing Tableau dashboards.
-    '
+    Chat History:
+        {context_chat_history}
+
+    Context:
+        {context_from_index}
+
+    User Query:
+        {prompt}
     """
-    print("Generating Chat Message...")
-    print(f"User Message:\n{prompt}\n")
-    chat_messages.append(prompt)
-    
+
     try:
-        response = genai.chat(
-          **chat_defaults,
-          context=context,
-          messages=chat_messages
+        response = genai.generate_text(
+            prompt=prompt_with_context,
+            safety_settings=[
+                {
+                    'category': genai.types.HarmCategory.HARM_CATEGORY_UNSPECIFIED,
+                    'threshold': genai.types.HarmBlockThreshold.BLOCK_NONE,
+                },
+            ]
         )
-        result = response.last
-        if result is None:
-            result = "Apologies but something went wrong. Please try again later."
-            chat_messages = chat_messages[:-1]
-        else:
-            chat_messages.append(result)
-            
+        result = response.result
+
     except Exception as e:
-        result = "Apologies but something went wrong. Please try again later."
-        chat_messages = chat_messages[:-1]
-        print(f"Exception {e} occured\n")
-        
-    chat_history = chat_messages
-    print(f"Bot Message:\n{result}\n")
+        result = "Seems something went wrong. Please try again later."
+        logger.error(f"Exception {e} occured\n")
+
+    chat_history.append(result)
+    logger.info(f"Bot Message:\n{result}\n")
     return result
 
 with gr.Blocks() as app:
@@ -74,11 +85,6 @@ with gr.Blocks() as app:
             placeholder="Hi Gerard, can you introduce yourself?",
             container=False, 
             scale=6)
-        # send = gr.Button(
-        #     value="", 
-        #     icon="./send-message.png", 
-        #     scale=1
-        # )
     clear = gr.Button("Clear")
 
     def user(user_message, history):
@@ -95,9 +101,6 @@ with gr.Blocks() as app:
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         bot, chatbot, chatbot
     )
-    # send.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-    #     bot, chatbot, chatbot
-    # )
     clear.click(clear_chat, None, chatbot, queue=False)
 
     gr.HTML("<p><center>Hosted on 🤗 Spaces. Powered by Google PaLM 🌴</center></p>")

requirements.txt

@@ -1,2 +1,3 @@
 google-generativeai
-langchain
+llama-index
+pinecone-client

src/llamaindex_palm.py

@@ -0,0 +1,171 @@
+import os
+import logging
+
+from typing import Any, List
+from pydantic import Extra
+
+import pinecone
+import google.generativeai as genai
+
+from llama_index import (
+    ServiceContext,
+    PromptHelper,
+    VectorStoreIndex
+)
+from llama_index.vector_stores import PineconeVectorStore
+from llama_index.storage.storage_context import StorageContext
+from llama_index.node_parser import SimpleNodeParser
+from llama_index.text_splitter import TokenTextSplitter
+from llama_index.embeddings.base import BaseEmbedding
+from llama_index.llms import (
+    CustomLLM,
+    CompletionResponse,
+    CompletionResponseGen,
+    LLMMetadata,
+)
+from llama_index.llms.base import llm_completion_callback
+
+class LlamaIndexPaLMEmbeddings(BaseEmbedding, extra=Extra.allow):
+    def __init__(
+        self,
+        model_name: str = 'models/embedding-gecko-001',
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._model_name = model_name
+
+    @classmethod
+    def class_name(cls) -> str:
+        return 'PaLMEmbeddings'
+
+    def gen_embeddings(self, text: str) -> List[float]:
+        return genai.generate_embeddings(self._model_name, text)
+
+    def _get_query_embedding(self, query: str) -> List[float]:
+        embeddings = self.gen_embeddings(query)
+        return embeddings['embedding']
+
+    def _get_text_embedding(self, text: str) -> List[float]:
+        embeddings = self.gen_embeddings(text)
+        return embeddings['embedding']
+
+    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        embeddings = [
+            self.gen_embeddings(text)['embedding'] for text in texts
+        ]
+        return embeddings
+    
+    async def _aget_query_embedding(self, query: str) -> List[float]:
+        return self._get_query_embedding(query)
+
+    async def _aget_text_embedding(self, text: str) -> List[float]:
+        return self._get_text_embedding(text)
+    
+class LlamaIndexPaLMText(CustomLLM, extra=Extra.allow):
+    def __init__(
+        self,
+        model_name: str = 'models/text-bison-001',
+        context_window: int = 8196,
+        num_output: int = 1024,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._model_name = model_name
+        self._context_window = context_window
+        self._num_output = num_output
+        
+    @property
+    def metadata(self) -> LLMMetadata:
+        """Get LLM metadata."""
+        return LLMMetadata(
+            context_window=self._context_window,
+            num_output=self._num_output,
+            model_name=self._model_name
+        )
+
+    def gen_texts(self, prompt):
+            logging.debug(f"prompt: {prompt}")
+            response = genai.generate_text(
+                model=self._model_name,
+                prompt=prompt,
+                safety_settings=[
+                    {
+                        'category': genai.types.HarmCategory.HARM_CATEGORY_UNSPECIFIED,
+                        'threshold': genai.types.HarmBlockThreshold.BLOCK_NONE,
+                    },
+                ]
+            )
+            logging.debug(f"response:\n{response}")
+            return response.candidates[0]['output']
+
+    @llm_completion_callback()
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        text = self.gen_texts(prompt)
+        return CompletionResponse(text=text)
+
+    @llm_completion_callback()
+    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+        raise NotImplementedError()
+    
+class LlamaIndexPaLM():
+    def __init__(
+        self,
+        emb_model: LlamaIndexPaLMEmbeddings = LlamaIndexPaLMEmbeddings(),
+        model: LlamaIndexPaLMText = LlamaIndexPaLMText()
+    ) -> None:
+        self.emb_model = emb_model
+        self.llm = model
+        
+        # Google Generative AI
+        genai.configure(api_key=os.environ['PALM_API_KEY'])
+
+        # Pinecone
+        pinecone.init(
+            api_key=os.environ['PINECONE_API_KEY'],
+            environment=os.getenv('PINECONE_ENV', 'us-west1-gcp-free')
+        )
+
+        # model metadata
+        CONTEXT_WINDOW = os.getenv('CONTEXT_WINDOW', 8196)
+        NUM_OUTPUT = os.getenv('NUM_OUTPUT', 1024)
+        TEXT_CHUNK_SIZE = os.getenv('TEXT_CHUNK_SIZE', 512)
+        TEXT_CHUNK_OVERLAP = os.getenv('TEXT_CHUNK_OVERLAP', 20)
+        TEXT_CHUNK_OVERLAP_RATIO = os.getenv('TEXT_CHUNK_OVERLAP_RATIO', 0.1)
+        TEXT_CHUNK_SIZE_LIMIT = os.getenv('TEXT_CHUNK_SIZE_LIMIT', None)
+
+        self.node_parser = SimpleNodeParser.from_defaults(
+            text_splitter=TokenTextSplitter(
+                chunk_size=TEXT_CHUNK_SIZE, chunk_overlap=TEXT_CHUNK_OVERLAP
+            )
+        )
+
+        self.prompt_helper = PromptHelper(
+            context_window=CONTEXT_WINDOW,
+            num_output=NUM_OUTPUT,
+            chunk_overlap_ratio=TEXT_CHUNK_OVERLAP_RATIO,
+            chunk_size_limit=TEXT_CHUNK_SIZE_LIMIT
+        )
+
+        self.service_context = ServiceContext.from_defaults(
+            llm=self.llm,
+            embed_model=self.embed_model,
+            node_parser=self.node_parser,
+            prompt_helper=self.prompt_helper,
+        )
+    
+    def set_index_from_pinecone(
+        self, 
+        index_name: str = 'experience'
+    ) -> None:
+        # Pinecone VectorStore
+        pinecone_index = pinecone.Index(index_name)
+        self.vector_store = PineconeVectorStore(pinecone_index=pinecone_index, add_sparse_vector=True)
+        self.pinecone_index = VectorStoreIndex.from_vector_store(self.vector_store, self.service_context)
+        return None
+    
+    def generate_response(
+        self,
+        query: str
+    ) -> str:
+        response = self.pinecone_index.as_query_engine().query(query)
+        return response.response