semsearch.py

import weaviate
from weaviate.connect import ConnectionParams
from weaviate.classes.init import AdditionalConfig, Timeout

from sentence_transformers import SentenceTransformer
from langchain_community.document_loaders import BSHTMLLoader
from pathlib import Path
from lxml import html
import logging
from semantic_text_splitter import HuggingFaceTextSplitter
from tokenizers import Tokenizer
import json
import os
import re
import logging

import llama_cpp
from llama_cpp import Llama

import streamlit as st
import subprocess


st.markdown("<h1 style='text-align: center; color: #666666;'>Vector Database RAG Proof of Concept</h1>", unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: #666666;'>V1</h6>", unsafe_allow_html=True)


weaviate_logger = logging.getLogger("httpx")
weaviate_logger.setLevel(logging.WARNING)

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def runStartup():
    result = subprocess.run(["bash", "startup.sh"], capture_output=True, text=True)
    logger(f"startup.sh stdout:  {result.stdout}")
    logger(f"startup.sh stderr:  {result.stderr}")
    logger(f"Return code: {result.returncode}")
logger("### Running startup.sh")
runStartup()

# Function to load the CSS file
def load_css(file_name):
    with open(file_name) as f:
        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)

# Load the custom CSS
load_css(".streamlit/main.css")

st.markdown("<h1 style='text-align: center; color: #666666;'>Vector Database RAG Proof of Concept</h1>", unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: #666666;'>V1</h6>", unsafe_allow_html=True)

######################################################################
# MAINLINE
#
logger.info("#### MAINLINE ENTERED.")

#pathString = "/Users/660565/KPSAllInOne/ProgramFilesX86/WebCopy/DownloadedWebSites/LLMPOC_HTML" 
pathString = "/app/inputDocs"
chunks = []
webpageDocNames = []
page_contentArray = []
webpageChunks = []
webpageTitles = []
webpageChunksDocNames = []

######################################################
# Connect to the Weaviate vector database.
logger.info("#### Create Weaviate db client connection.")
client = weaviate.WeaviateClient(
    connection_params=ConnectionParams.from_params(
        http_host="localhost",
        http_port="8080",
        http_secure=False,
        grpc_host="localhost",
        grpc_port="50051",
        grpc_secure=False
    ),
    additional_config=AdditionalConfig(
        timeout=Timeout(init=60, query=1800, insert=1800),  # Values in seconds
    )
)
client.connect() 


#######################################################
# Read each text input file, parse it into a document,
# chunk it, collect chunks and document name.
logger.info("#### Read and chunk input text files.")
if not client.collections.exists("Documents") or not client.collections.exists("Documents") :
    for filename in os.listdir(pathString):
        logger.info(filename)
        path = Path(pathString + "/" + filename)
        filename = filename.rstrip(".html")
        webpageDocNames.append(filename)
        htmlLoader = BSHTMLLoader(path,"utf-8") 
        htmlData = htmlLoader.load()

        title   = htmlData[0].metadata['title']
        page_content = htmlData[0].page_content
        
        # Clean data. Remove multiple newlines, etc.
        page_content = re.sub(r'\n+', '\n',page_content)
        
        page_contentArray.append(page_content);
        webpageTitles.append(title)
        max_tokens = 1000
        tokenizer = Tokenizer.from_pretrained("bert-base-uncased")
        logger.debug(f"### tokenizer: {tokenizer}")
        splitter = HuggingFaceTextSplitter(tokenizer, trim_chunks=True)
        chunksOnePage = splitter.chunks(page_content, chunk_capacity=50)

        chunks = []
        for chnk in chunksOnePage:
            logger.debug(f"#### chnk in file: {chnk}")
            chunks.append(chnk)
            logger.debug(f"chunks: {chunks}")
        webpageChunks.append(chunks)
        webpageChunksDocNames.append(filename + "Chunks")

        logger.debug(f"### filename, title: {filename}, {title}")
        logger.debug(f"### webpageDocNames: {webpageDocNames}")



######################################################
# Create database webpage and chunks collections.
#wpCollection = createWebpageCollection()
#wpChunkCollection = createChunksCollection()
logger.info("#### createWebpageCollection() entered.")
if not client.collections.exists("Documents"):
    #client.collections.delete("Documents")
    class_obj = {
      "class": "Documents",
      "description": "For first attempt at loading a Weviate database.",
      "vectorizer": "text2vec-transformers",
      "moduleConfig": {
        "text2vec-transformers": {
        "vectorizeClassName": False
        }
      },
      "vectorIndexType": "hnsw",  
      "vectorIndexConfig": {
          "distance": "cosine",
      },
      "properties": [
            {
                "name": "title",
                "dataType": ["text"],
                "description": "HTML doc title.",
                "vectorizer": "text2vec-transformers",
                "moduleConfig": {
                    "text2vec-transformers": {
                        "vectorizePropertyName": True,
                        "skip": False,
                        "tokenization": "lowercase"
                    }
                },
                "invertedIndexConfig": {
                    "bm25": {
                        "b": 0.75,
                        "k1": 1.2
                    },
                }
            },
            {
                "name": "content",
                "dataType": ["text"],
                "description": "HTML page content.",
                "moduleConfig": {
                    "text2vec-transformers": {
                        "vectorizePropertyName": True,
                        "tokenization": "whitespace"
                    }
                }
            }
        ]  
    }
    wpCollection = client.collections.create_from_dict(class_obj)

logger.info("#### createChunksCollection() entered.")
if not client.collections.exists("Chunks"):
    #client.collections.delete("Chunks")
    class_obj = {
        "class": "Chunks",
        "description": "Collection for document chunks.",
        "vectorizer": "text2vec-transformers",
        "moduleConfig": {
            "text2vec-transformers": {
                "vectorizeClassName": True
            }
        },
        "vectorIndexType": "hnsw",  
        "vectorIndexConfig": {
            "distance": "cosine",
        },
        "properties": [
            {
                "name": "chunk",
                "dataType": ["text"],
                "description": "Single webpage chunk.",
                "vectorizer": "text2vec-transformers",
                "moduleConfig": {
                    "text2vec-transformers": {
                      "vectorizePropertyName": False,
                      "skip": False,
                      "tokenization": "lowercase"
                    }
                }
            },
            {
                "name": "chunk_index",
                "dataType": ["int"]
            },
            {
                "name": "webpage",
                "dataType": ["Documents"],
                "description": "Webpage content chunks.",
            
                "invertedIndexConfig": {
                    "bm25": {
                        "b": 0.75,
                        "k1": 1.2
                    }
                }  
            }
        ]
    }
    wpChunkCollection = client.collections.create_from_dict(class_obj)


###########################################################
# Create document and chunks objects in the database.
if not client.collections.exists("Documents") :
    logger.info("#### Create page/doc db objects.")
    for i, className in enumerate(webpageDocNames):
        title = webpageTitles[i]
        logger.debug(f"## className, title: {className}, {title}")
        # Create Webpage Object
        page_content = page_contentArray[i]
        # Insert the document.
        wpCollectionObj_uuid = wpCollection.data.insert(
          {
            "name": className,
            "title": title,
            "content": page_content
          }
        )
    
if not client.collections.exists("Chunks") :
    logger.info("#### Create chunk db objects.")
    # Insert the chunks for the document.
    for i2, chunk in enumerate(webpageChunks[i]):
        chunk_uuid = wpChunkCollection.data.insert(
          {
            "title": title,
            "chunk": chunk,
            "chunk_index": i2,
            "references":
            {
              "webpage": wpCollectionObj_uuid
            }
          }
        )


#################################################################
# Initialize the LLM.
model_path = "/app/llama-2-7b-chat.Q4_0.gguf"
llm = Llama(model_path,
            #*,
            n_gpu_layers=0,
            split_mode=llama_cpp.LLAMA_SPLIT_MODE_LAYER,
            main_gpu=0,
            tensor_split=None,
            vocab_only=False,
            use_mmap=True,
            use_mlock=False,
            kv_overrides=None,
            seed=llama_cpp.LLAMA_DEFAULT_SEED,
            n_ctx=512,
            n_batch=512,
            n_threads=8,
            n_threads_batch=16,
            rope_scaling_type=llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED,
            pooling_type=llama_cpp.LLAMA_POOLING_TYPE_UNSPECIFIED,
            rope_freq_base=0.0,
            rope_freq_scale=0.0,
            yarn_ext_factor=-1.0,
            yarn_attn_factor=1.0,
            yarn_beta_fast=32.0,
            yarn_beta_slow=1.0,
            yarn_orig_ctx=0,
            logits_all=False,
            embedding=False,
            offload_kqv=True,
            last_n_tokens_size=64,
            lora_base=None,
            lora_scale=1.0,
            lora_path=None,
            numa=False,
            chat_format=None,
            chat_handler=None,
            draft_model=None,
            tokenizer=None,
            type_k=None,
            type_v=None,
            verbose=True
           )


def getRagData(promptText):
    ###############################################################################
    # Initial the the sentence transformer and encode the query prompt.
    logger.info(f"#### Encode text query prompt to create vectors. {text}")
    model = SentenceTransformer('/app/multi-qa-MiniLM-L6-cos-v1')

    vector = model.encode(promptText)
    vectorList = []

    logger.debug("#### Print vectors.")
    for vec in vector:
        vectorList.append(vec)
    logger.debug(f"vectorList: {vectorList[2]}")

    # Fetch chunks and print chunks.
    logger.info("#### Retrieve semchunks from db using vectors from prompt.")
    semChunks = wpChunkCollection.query.near_vector(
        near_vector=vectorList,
        distance=0.7,
        limit=3
    )
    logger.debug(f"### semChunks[0]: {semChunks}")

    # Print chunks, corresponding document and document title.
    ragData = ""
    logger.info("#### Print individual retrieved chunks.")
    for chunk in enumerate(semChunks.objects):
        logger.info(f"#### chunk: {chunk}")
        ragData = ragData + "\n" + chunk[0]
        webpage_uuid = chunk[1].properties['references']['webpage']
        logger.info(f"webpage_uuid: {webpage_uuid}")
        wpFromChunk = wpCollection.query.fetch_object_by_id(webpage_uuid)
        logger.info(f"### wpFromChunk title: {wpFromChunk.properties['title']}")
    #collection = client.collections.get("Chunks")
    return  ragData


# Display UI
col1, col2 = st.columns(2)

with col1:
    if "sysTA" not in st.session_state:
        st.session_state.sysTA = st.text_area(label="sysTA",value="fdsaf fsdafdsa")
    elif "sysTAtext" in st.session_state:
        st.session_state.sysTA = st.text_area(label="sysTA",value=st.session_state.sysTAtext)
    else:
        st.session_state.sysTA = st.text_area(label="sysTA",value=st.session_state.sysTA)
        
    if "userpTA" not in st.session_state:
        userTextArea = st.text_area(label="userpTA",value="fdsaf fsdafdsa")
    elif "userpTAtext" in st.session_state:
        st.session_state.userpTA = st.text_area(label="userpTA",value=st.session_state.userpTAtext)
    else:
        st.session_state.userpTA = st.text_area(label="userpTA",value=st.session_state.userpTA)

with col2:
    if "ragpTA" not in st.session_state:
        ragPromptTextArea = st.text_area(label="ragpTA",value="fdsaf fsdafdsa")
    elif "ragpTAtext" in st.session_state:
        st.session_state.ragpTA = st.text_area(label="ragpTA",value=st.session_state.ragpTAtext)
    else:
        st.session_state.ragTA = st.text_area(label="ragTA",value=st.session_state.ragTA)

    if "rspTA" not in st.session_state:
        responseTextArea = st.text_area(label="rspTA",value="fdsaf fsdafdsa")
    elif "rspTAtext" in st.session_state:
        st.session_state.rspTA = st.text_area(label="rspTA",value=st.session_state.rspTAtext)
    else:
        st.session_state.rspTA = st.text_area(label="rspTA",value=st.session_state.rspTA)

def runLLM(prompt):
    max_tokens = 1000
    temperature = 0.3
    top_p = 0.1
    echo = True
    stop = ["Q", "\n"]

    modelOutput = llm(
       prompt,
       max_tokens=max_tokens,
       temperature=temperature,
       top_p=top_p,
       echo=echo,
       stop=stop,
    )
    result = modelOutput["choices"][0]["text"].strip()
    return(result)

def setPrompt(pprompt,ragFlag):
    print("\n### setPrompt() entered. ragFlag: ",ragFlag)
    if ragFlag:
        ragPrompt = getRagData(pprompt)
        userPrompt = pprompt + "\n" + ragPrompt
        prompt = userPrompt
        userPrompt = "Using this information: " + ragPrompt \
             + "process the following statement or question and produce a a response" \
             + intialPrompt
    else:
        userPrompt = pprompt
    #prompt = f""" <s> [INST] <<SYS>> {systemTextArea.value} </SYS>> Q: {userPrompt} A: [/INST]"""
    return userPrompt


def on_submitButton_clicked(b):
    logger.debug("\n### on_submitButton_clicked")
    st.session_state.sysTAtext = st.session_state.sysTA
    logger.info(f"sysTAtext: {st.session_state.sysTAtext}")
    
    st.session_state.userpTAtext = setPrompt("","")
    st.session_state.userpTA = st.session_state.userpTAtext
    logger.info(f"userpTAtext: {st.session_state.userpTAtext}")
    
    st.session_state.rspTAtext = runLLM(st.session_state.userpTAtext)
    st.session_state.rspTA = st.session_state.rspTAtext
    logger.info(f"rspTAtext: {st.session_state.rspTAtext}")


with st.sidebar:
    st.selectRag    = st.checkbox("Enable Query With RAG",value=False,key="selectRag",help=None,on_change=None,args=None,kwargs=None,disabled=False,label_visibility="visible")
    st.submitButton = st.button("Run LLM Query",key=None,help=None,on_click=on_submitButton_clicked,args=None,kwargs=None,type="secondary",disabled=False,use_container_width=False)

#logger.info("#### Closing client db connection.")
#client.close()
#logger.info("#### Program terminating.")