|
import os |
|
import re |
|
import lz4 |
|
import json |
|
import time |
|
import uuid |
|
import torch |
|
import base64 |
|
import asyncio |
|
import msgpack |
|
import validators |
|
import numpy as np |
|
import pandas as pd |
|
import streamlit as st |
|
|
|
from datetime import datetime |
|
from streamlit import _bottom |
|
from pydantic import BaseModel |
|
from streamlit_pills import pills |
|
from dense_embed import embed_text |
|
from ppt_chunker import ppt_chunker |
|
from streamlit_navigation_bar import st_navbar |
|
from vllm.sampling_params import SamplingParams |
|
from outlines.fsm.json_schema import build_regex_from_schema |
|
from unstructured.nlp.tokenize import download_nltk_packages |
|
|
|
from typing import ( |
|
List, |
|
Dict |
|
) |
|
from fastembed import ( |
|
SparseTextEmbedding, |
|
SparseEmbedding |
|
) |
|
from infinity_emb import ( |
|
AsyncEngineArray, |
|
EngineArgs, |
|
AsyncEmbeddingEngine |
|
) |
|
from search_strategies import ( |
|
QdrantClient, |
|
query_keywords_search, |
|
query_hybrid_search |
|
) |
|
from scipy.sparse import ( |
|
csr_matrix, |
|
save_npz, |
|
load_npz, |
|
vstack |
|
) |
|
from prompts import ( |
|
outlines, |
|
build_prompt_conv, |
|
route_llm, |
|
open_query_prompt, |
|
question_type_prompt, |
|
idk, |
|
self_knowledge, |
|
answer_with_context |
|
) |
|
from qdrant_client.models import ( |
|
NamedSparseVector, |
|
NamedVector, |
|
SparseVector, |
|
PointStruct, |
|
ScoredPoint, |
|
SearchRequest, |
|
Modifier, |
|
OptimizersConfigDiff, |
|
HnswConfigDiff, |
|
Distance, |
|
VectorParams, |
|
SparseVectorParams, |
|
SparseIndexParams, |
|
Batch, |
|
Filter, |
|
FieldCondition, |
|
Datatype, |
|
MatchValue |
|
) |
|
|
|
class Question(BaseModel): |
|
answer: str |
|
schema = json.dumps(Question.model_json_schema()) |
|
|
|
icon_to_types = { |
|
'ppt':('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADAAAAAwCAYAAABXAvmHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAC4ElEQVR4nO2YS2gTQRzGF/VgPVuPpsXSQ3eshdbXqVDMEtlN1UN8HqTZYouCVJTsiIfgTcWbD0o0mp3GKq1oC55E2AbbKn3dakISCVjwUNtLlZ1a24xs0oqG2uxmJ2si+8F32NnL95v/fx67DGPLli1b/5UUkauJtDsvDHm5XkXkphSvc17xcksqgksYSXMYwSksS71Yhudx6MouphSk+Ju3RLzcGUV0jg6JHFnPGMH1LcMRVZZOkz7P5n8TXuRcisjF/xY8LwDKWpVhDIcgZ1nwiXPCNkXkHuYLrhcAr4EgKUD6LlUUNfxIp3OHIjon9YY3AoCzHl8IXq0sWvgh0RkzEr4AAKK1FHWIbNsYm/lCAfBqJchj/1ZqAEZ6nhIAURHsNhQy0OToCjQ5vj1ochAzDu6vJgOte00DYATJYsh32AiA6fC/Q9AAUJEU1X1O0Aq/ZhoAOAPhO1nWABhJw2UNoMowjdG16oIBcrWy/IPMx6Pk9eV2iyoACQ75OkwDxF4+JdEXT8jMaCTznF5ZJoNtR60BkKWwaYDgwZpfY/FXzzNj0/3IGgAEJ6gCjN29ma3KwDOLAKRZ0wBhrpGglnoy2HaMLHyeyYy98XVaAqAiuGgaIFcf+ns2XMRJAdD0ommA8Xu3yNidG+Td7etk4Gxr3m2ULgA7S3UN6DFlgIlyB+gpcwBQ+EFWqGmFTwggnTqyp6p8AXjwNm/4kgZwg+N6Ab7SCv9oXxUdAB5ME4/eD5rGnRdpQGjhIy21dADcdS0MDSUFEKC8qxAdi/c+Q0ufPAcqEjw7biHA+1Szg95vFU1xV0NlgmdjRZ95no3GhNrtVMP/CQHGijnzcVdDcX4t5rRTdzF6PkW7bTbSR373Ia3cpsPzYJrabmNU2h6ddNedSvBgWDvyDcx2WjthEwJ7gviZTUwpKOaur9YuXQmBDSd5djLJg7kkD76v+ot2Jc68E0CHrruNLVu2bDGlrJ8c/urSuEn7XgAAAABJRU5ErkJggg==', |
|
'Powerpoint'), |
|
'pptx':('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADAAAAAwCAYAAABXAvmHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAC4ElEQVR4nO2YS2gTQRzGF/VgPVuPpsXSQ3eshdbXqVDMEtlN1UN8HqTZYouCVJTsiIfgTcWbD0o0mp3GKq1oC55E2AbbKn3dakISCVjwUNtLlZ1a24xs0oqG2uxmJ2si+8F32NnL95v/fx67DGPLli1b/5UUkauJtDsvDHm5XkXkphSvc17xcksqgksYSXMYwSksS71Yhudx6MouphSk+Ju3RLzcGUV0jg6JHFnPGMH1LcMRVZZOkz7P5n8TXuRcisjF/xY8LwDKWpVhDIcgZ1nwiXPCNkXkHuYLrhcAr4EgKUD6LlUUNfxIp3OHIjon9YY3AoCzHl8IXq0sWvgh0RkzEr4AAKK1FHWIbNsYm/lCAfBqJchj/1ZqAEZ6nhIAURHsNhQy0OToCjQ5vj1ochAzDu6vJgOte00DYATJYsh32AiA6fC/Q9AAUJEU1X1O0Aq/ZhoAOAPhO1nWABhJw2UNoMowjdG16oIBcrWy/IPMx6Pk9eV2iyoACQ75OkwDxF4+JdEXT8jMaCTznF5ZJoNtR60BkKWwaYDgwZpfY/FXzzNj0/3IGgAEJ6gCjN29ma3KwDOLAKRZ0wBhrpGglnoy2HaMLHyeyYy98XVaAqAiuGgaIFcf+ns2XMRJAdD0ommA8Xu3yNidG+Td7etk4Gxr3m2ULgA7S3UN6DFlgIlyB+gpcwBQ+EFWqGmFTwggnTqyp6p8AXjwNm/4kgZwg+N6Ab7SCv9oXxUdAB5ME4/eD5rGnRdpQGjhIy21dADcdS0MDSUFEKC8qxAdi/c+Q0ufPAcqEjw7biHA+1Szg95vFU1xV0NlgmdjRZ95no3GhNrtVMP/CQHGijnzcVdDcX4t5rRTdzF6PkW7bTbSR373Ia3cpsPzYJrabmNU2h6ddNedSvBgWDvyDcx2WjthEwJ7gviZTUwpKOaur9YuXQmBDSd5djLJg7kkD76v+ot2Jc68E0CHrruNLVu2bDGlrJ8c/urSuEn7XgAAAABJRU5ErkJggg==', |
|
'Powerpoint'), |
|
'txt':('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAYAAACqaXHeAAAACXBIWXMAAAsTAAALEwEAmpwYAAAEA0lEQVR4nO2bW08TQRTH+8QDH8DPpF+EGLt+AENI44OC7u6L+mAxYQ0xLdZoDE8kCpr4oEYUGu4kEOgFRWm5FLbd3Zk9Ztpus132WmZ3WtqTnGQzLbP7/82ZM2dmaSw2sIENbGAdWHIkOSxw0pgQlzYFTlJFTgIGXhQ4aSKRyAzFojT+9vQNMT71i5HoS04gRCpe4KZWyY2TiUzrIdTj3cCOzooA1ZO64/ND23azm79DfD+73IqEyMSLcWmF3HQykYG5ldOOAVxVvHayB7pcbt2fifgPu0pHAGiJJ59FAoB3EN8JAJriIwHAu4gPCoC2+NAB8B7i/QDA1ePQxIcKgPch3gsAkv8BMQIhDPGhAeB9incDgOQStJlWDSyegHMTHwqA5Ehy2ChyXtx/A3OrZ47ivSIAGxCUi8DitcpB/U915dxRfCgABE4aM4ocL/Hvv+TrN39+L+Wc8FS5I/E6VhvwdB1AvYgQQFzaDFqKLryep5btUbUMtoaUyCJA9SucjPznzDxUj3boLnVKBUDHpgiQo4sAkWFt35bwyNwn5iK+qwCgMNZ5reYqvmsAoJCKHD/OHABiKJ4pgK3vi5AcTTE/DPHtcSnPcy9vUQMwOZpmLyqgC3elAhUA6KzY6vTBj1pPuO+pInoAMOZ2XwJApsTWdwCQJdvTBPD4Zw0mFrsYALJZ6vwAyFea5ayN5Sq4/p2HizWoqDocynobBHJ9cK7X+wjSH3UAyGGd9wMg5/OBiXhiC3mt1Uauif2V9cD9UQOAXIqcoFPAMLvPpjcaW18VAzzLKvAkq4DS1JreVgP3RwUA8qjwaAIgvlZqKF4vYVgvN6/LzqMaOgDwKG9pAzCPuhENT7MKWwDYpbanDcA87635gAkA7LGxCQPAp0IXAVA9dnW0AZBwJ2FvGJkOZFowBYBctrS0AWw0E9+aKQmSa2YAkMd+niaA1FZjGVRQY9TNCfHVJqNlEDwOM2gCIBUfsY+5y4VQwaHAiRwAtpzkBAVAKjWnau3rb1QP+3FLKbx9jOHbHxS4P+oAsM0x1rXfDYLHGV5fAMDmWkAugXaau/SfGdcWAG4Tf0Re0QBgDaB22h8A1Law3wcgLyvJK6vm1JgcnYn0QJPZoajWOrdvjLzhu0tbPQVB4KS8eEe6GQiAdsWXFl5uTbCduO/Q9mtis8NeEB8qAL0HxIcKAHpAfM8BwJTF9xQAHIL4SACk+VmYEWfbxARtM4tPPXoLaf5dm4irtIUOQLSJiKBtdg9Mu20AgJaJlpEkIWwNbT9tJOxJuFpDlnZbVyZBHFLCiyoJFkmHhfWdrhe/t7xE/yczAidNsN64BPb41Dg1AIlEZqgJoR4JXe5sfjY3sIENLHZd7D/x1k4dCUv1GwAAAABJRU5ErkJggg==', |
|
'Txt'), |
|
'doc':('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADAAAAAwCAYAAABXAvmHAAAACXBIWXMAAAsTAAALEwEAmpwYAAACOElEQVR4nO2ZQWgTQRSGh+DFq70VD2lvQuKpV714cGtLu3ioiiJeWr0oaSnSm0KhKnabJiKYYhHRQxORBmLqoYVKqaGGWmizAYsaCBhDpdCkqatpap68l3QO4kHJsBnN/PDDzts34f+YN3sJY0pKSkr/pFz31q64/Ym822+CPU7kXX7zsjAAe8ObZJfPzAkEsDe8u2oFsC91An41Qv/rCGmGh2nGDtMMqMWOznFoHpipC8BOreE5RMc4uIdNoWZ/AAAirQCGG/kEjg8GARVZ+shrR3ofwb5ODD3j9Wg8RTXcIw3AwS4ffN/dgy85i9f6fLMc4MaTGK9vbH2lXtwjDQDTDHidzFDYlksPaf14LgnFUgVqdiVNNXyHwl6pRohpBtwJxSncmZEXtE5lc7D07jNMx95DwdqFAx1eOHc7Sj3YKx1A180whTOeL0Pz+QA9e6ffwvXJBXpuu/qU1ijslQ6gqec+lMsAi8kMnQKqZyQCxwan6NnzYJ5GB3uw99f9Rz1vhJr9LQA6md6Eb8USBGZWKfThCxN0WfEuhGMfwCqWqOd3e6UACFSDb1tF+LRZ4PX4ehZKez/oHfZIC3Dx7kv+6QwtrPO6L7zC69gjLUBr9TOJ6g+84vWztyp3AoU90gKwGqwAPI13AqMFUeEdp7x1ADhpXBMB4WgfKx/qnLRaT4e2RJqJUosehHqYKYCq1AnoaoQafIScejBvd3hn95S4PziceqgPf9DO8E492CsMQElJSYmJ0E+635eFCoKREwAAAABJRU5ErkJggg==', |
|
'Microsoft Word'), |
|
'docx':('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADAAAAAwCAYAAABXAvmHAAAACXBIWXMAAAsTAAALEwEAmpwYAAACOElEQVR4nO2ZQWgTQRSGh+DFq70VD2lvQuKpV714cGtLu3ioiiJeWr0oaSnSm0KhKnabJiKYYhHRQxORBmLqoYVKqaGGWmizAYsaCBhDpdCkqatpap68l3QO4kHJsBnN/PDDzts34f+YN3sJY0pKSkr/pFz31q64/Ym822+CPU7kXX7zsjAAe8ObZJfPzAkEsDe8u2oFsC91An41Qv/rCGmGh2nGDtMMqMWOznFoHpipC8BOreE5RMc4uIdNoWZ/AAAirQCGG/kEjg8GARVZ+shrR3ofwb5ODD3j9Wg8RTXcIw3AwS4ffN/dgy85i9f6fLMc4MaTGK9vbH2lXtwjDQDTDHidzFDYlksPaf14LgnFUgVqdiVNNXyHwl6pRohpBtwJxSncmZEXtE5lc7D07jNMx95DwdqFAx1eOHc7Sj3YKx1A180whTOeL0Pz+QA9e6ffwvXJBXpuu/qU1ijslQ6gqec+lMsAi8kMnQKqZyQCxwan6NnzYJ5GB3uw99f9Rz1vhJr9LQA6md6Eb8USBGZWKfThCxN0WfEuhGMfwCqWqOd3e6UACFSDb1tF+LRZ4PX4ehZKez/oHfZIC3Dx7kv+6QwtrPO6L7zC69gjLUBr9TOJ6g+84vWztyp3AoU90gKwGqwAPI13AqMFUeEdp7x1ADhpXBMB4WgfKx/qnLRaT4e2RJqJUosehHqYKYCq1AnoaoQafIScejBvd3hn95S4PziceqgPf9DO8E492CsMQElJSYmJ0E+635eFCoKREwAAAABJRU5ErkJggg==', |
|
'Microsoft Word'), |
|
'xslx':('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADAAAAAwCAYAAABXAvmHAAAACXBIWXMAAAsTAAALEwEAmpwYAAADGUlEQVR4nO2Za0hTYRzG9z0vu59JpK3vedvmTHOmM1OnE0wqbZWJjjTNS2nz7jZvFVmJUTOCchF0+RaBiAThQkgQKSKKbjohIqMIKlG3J3aiVXicw3PaMToPPF/eT7/fe/7vew4cHo8LFy5c/qlssht1cnvpjNxuxGqqcJhX1dhRs1M5as6kLbBx0OhcLTwdAYVHwtE+TVuADjxdAYXDDE5Azj0BMzdCtLLhTCH8aXifAREXiyC/WhrYQ0z06HVET+4M0a0H0ZUDojMb0g4dpNYsSC1ZkJgzIWnPgKRtB8Rt6RC3pkPcsh2i5jSImrQQNqZC2JACUYMWMose4f2GPyT+vkC33kkXXmhKgeD4NkhMqVB1FyHJVoHkKzVk6V4CPD8EGIEX1CdDUK+BqusAEs+XB1Cgizl4fl3SD4H+sgAKdDIHzz/GgoCUAn69NRsvZ2fgSeuQzQsffaoQcwvz+DY/h5gTBUvg+Ue3siBgpd75/dfbSYHZL58QYckhd/7W5Ai5Zh26RAkfWpvIgoBl+bEZef6QBO4YvgxNXwlcbjdevHdC1phGCR9akxB4AYmPmU84V4z5xQV8/PoZY68fwe12Q2+rXhY+tHoLCwLtvg/sBcdt/MyNiWGf8CFV8SwItPm+bezjd70CNyeGfcKHHImHOCcS0l0KyAxqsvzmVFrlrRRf8DkDNeTYjE89wYNXk6RE3kDtsvAhlWoWBFqp4cOa0/Hs3RsSWm+rgrbPSMpMfXiLMJOWEj64Mo4FgRbql9Tpe9dI+OGnY96xufP4PrnWOzJICR9coQq8gIgCPqH3IHn7uNwuaM4We2de2VOABdci2cST+5bABx1WsiDQtPLnga+Z/x0+qJwFASGD8EFlChYEGpiDX3coNvACAlOKkyn4oJJoFgTqNZn8Oo2TCXjBzkiIc6NA7P4lQLc8uhHpNsPfivVRkObHgChQrR0BSV4M/KnU0/xYELuVkO2NWzsCxB4l/K5n5xmElzEhwCSMjBP4H58AYVA72YInDGr6PzhkhepMNiQIg3paZojPoC3AhQsXLjwm8x3YSSmFlSW/AQAAAABJRU5ErkJggg==', |
|
'Excel') |
|
} |
|
|
|
def generate_answer(query: str, |
|
client: QdrantClient, |
|
collection_name: str, |
|
dense_model: AsyncEmbeddingEngine, |
|
sparse_model: SparseTextEmbedding, |
|
past_messages: str, |
|
search_strategy, |
|
documents_only: bool, |
|
gen_text, |
|
gen_context_choice, |
|
gen_question_choice |
|
): |
|
sparse_embeddings = list(sparse_model.query_embed(query))[0].as_object() |
|
|
|
s = time.time() |
|
|
|
if search_strategy == 'Exact Search': |
|
scored_point = query_keywords_search(query, client, collection_name, sparse_embeddings).points[0] |
|
|
|
text = scored_point.payload['text'] |
|
metadata = scored_point.payload['metadata']['url'] |
|
|
|
answer = f"{text}\n\n\nSource :\n\n{metadata}" |
|
else: |
|
dense_embeddings, tokens_count = asyncio.run(embed_text(dense_model[0], query)) |
|
|
|
scored_points = query_hybrid_search(query, client, collection_name, dense_embeddings, sparse_embeddings).points |
|
|
|
docs = [(scored_point.payload['text'], scored_point.payload['metadata']) for scored_point in scored_points] |
|
contents, metadatas = [list(t) for t in zip(*docs)] |
|
|
|
context = "\n".join(contents) |
|
|
|
prompt = route_llm(context, query) |
|
action = gen_context_choice(prompt, max_tokens=2, sampling_params=SamplingParams(temperature=0)) |
|
print(f'Choice: {action}') |
|
|
|
if action == 'Yes': |
|
filtered_metadatas = { |
|
value |
|
for metadata in metadatas |
|
if 'url' in metadata |
|
for value in [metadata['url']] |
|
} |
|
result_metadatas = "\n\n".join(f'{value}' for value in filtered_metadatas) |
|
|
|
prompt = answer_with_context(context, query) |
|
answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0)))['answer'] |
|
answer = f"{answer}\n\n\nSource(s) :\n\n{result_metadatas}" |
|
|
|
if documents_only == False: |
|
answer = f'Documents Based :\n\n{answer}' |
|
else: |
|
prompt = question_type_prompt(query) |
|
action = gen_question_choice(prompt, max_tokens=3, sampling_params=SamplingParams(temperature=0)) |
|
print(f'Choice 2: {action}') |
|
|
|
if action == 'General Question': |
|
prompt = open_query_prompt(past_messages, query) |
|
answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer'] |
|
else: |
|
if documents_only == True: |
|
prompt = idk(query) |
|
answer = json.loads(gen_text(prompt, max_tokens=128, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer'] |
|
else: |
|
prompt = self_knowledge(query) |
|
answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer'] |
|
answer = f'Internal Knowledge :\n\n{answer}' |
|
|
|
torch.cuda.empty_cache() |
|
|
|
e = time.time() |
|
f = e - s |
|
print(f'SEARCH TIME : {f}') |
|
|
|
return answer |
|
|
|
def collect_files(directory, pattern): |
|
array = [] |
|
|
|
for filename in os.listdir(directory): |
|
name, extension = filename.rsplit('.', 1) |
|
if name.endswith(pattern): |
|
print(f'Filename : {filename}') |
|
if filename.endswith('.msgpack'): |
|
with open(os.path.join(directory, filename), "rb") as data_file_payload: |
|
decompressed_payload = data_file_payload.read() |
|
array.extend(msgpack.unpackb(decompressed_payload, raw=False)) |
|
elif filename.endswith('.npz') and ('_dense' in filename): |
|
array.extend(list(np.load(os.path.join(directory, filename)).values())) |
|
elif filename.endswith('.npz') and ('_sparse' in filename): |
|
sparse_embeddings = [] |
|
loaded_sparse_matrix = load_npz(os.path.join(directory, filename)) |
|
|
|
for i in range(loaded_sparse_matrix.shape[0]): |
|
row = loaded_sparse_matrix.getrow(i) |
|
values = row.data.tolist() |
|
indices = row.indices.tolist() |
|
embedding = SparseVector(indices=indices, values=values) |
|
sparse_embeddings.append(embedding) |
|
array.extend(sparse_embeddings) |
|
elif filename.endswith('.npy'): |
|
ids_list = np.load(os.path.join(directory, filename), allow_pickle=True).tolist() |
|
array.extend(ids_list) |
|
|
|
return array |
|
|
|
@st.cache_resource(show_spinner=False) |
|
def load_models_and_documents(): |
|
container = st.empty() |
|
|
|
with container.status("Load AI Models and Prepare Documents...", expanded=True) as status: |
|
st.write('Downloading and Loading MixedBread Mxbai Dense Embedding Model with Infinity and PyTorch as backend...') |
|
|
|
dense_model = AsyncEngineArray.from_args( |
|
[ |
|
EngineArgs( |
|
model_name_or_path='intfloat/multilingual-e5-large-instruct', |
|
engine='torch', |
|
device='cpu', |
|
compile=False, |
|
embedding_dtype='float32', |
|
dtype='float16', |
|
pooling_method='cls', |
|
lengths_via_tokenize=True |
|
) |
|
] |
|
) |
|
|
|
st.write('Downloading and Loading Qdrant BM42 Sparse Embedding Model under ONNX using the CPU...') |
|
|
|
sparse_model = SparseTextEmbedding( |
|
'Qdrant/bm42-all-minilm-l6-v2-attentions', |
|
cache_dir=os.getenv('HF_HOME'), |
|
providers=['CPUExecutionProvider'] |
|
) |
|
|
|
st.write('Downloading and Loading Mistral Nemo quantized with GPTQ and using Outlines + vLLM Engine as backend...') |
|
|
|
llm = outlines.models.vllm( |
|
model_name="shuyuej/Mistral-Nemo-Instruct-2407-GPTQ", |
|
tensor_parallel_size=1, |
|
enforce_eager=True, |
|
gpu_memory_utilization=1, |
|
max_model_len=8192, |
|
dtype=torch.float16, |
|
max_num_seqs=64, |
|
quantization="gptq" |
|
) |
|
|
|
regex = build_regex_from_schema(schema, r"[\n ]?") |
|
gen_text = outlines.generate.regex(llm, regex) |
|
|
|
gen_context_choice = outlines.generate.choice(llm, choices=['Yes', 'No']) |
|
gen_question_choice = outlines.generate.choice(llm, choices=['Domain-Specific Question', 'General Question']) |
|
|
|
st.write('Downloading NLTK Packages...') |
|
|
|
download_nltk_packages() |
|
|
|
st.write('Creating Collection for our Qdrant Vector Database in RAM memory...') |
|
|
|
client = QdrantClient(path=os.path.join(os.getenv('HOME'), 'database_qdrant.sqlite')) |
|
collection_name = 'collection_demo' |
|
|
|
print(f'Client exist ? : {client.collection_exists(collection_name)}') |
|
|
|
client.create_collection( |
|
collection_name, |
|
{ |
|
'text-dense': VectorParams( |
|
size=1024, |
|
distance=Distance.COSINE, |
|
datatype=Datatype.FLOAT16, |
|
on_disk=False |
|
) |
|
}, |
|
{ |
|
'text-sparse': SparseVectorParams( |
|
index=SparseIndexParams( |
|
on_disk=False |
|
), |
|
modifier=Modifier.IDF |
|
), |
|
'title-sparse': SparseVectorParams( |
|
index=SparseIndexParams( |
|
on_disk=False |
|
), |
|
modifier=Modifier.IDF |
|
) |
|
}, |
|
2, |
|
optimizers_config=OptimizersConfigDiff( |
|
indexing_threshold=0, |
|
default_segment_number=4 |
|
), |
|
hnsw_config=HnswConfigDiff( |
|
on_disk=False, |
|
m=32, |
|
ef_construct=200 |
|
) |
|
) |
|
|
|
client.create_payload_index( |
|
collection_name=collection_name, |
|
field_name="url", |
|
field_schema="keyword" |
|
) |
|
|
|
name = 'action_rpg' |
|
embeddings_path = os.path.join(os.getenv('HF_HOME'), 'embeddings') |
|
|
|
payload_path = os.path.join(embeddings_path, name + '_payload.msgpack') |
|
payload_titles_path = os.path.join(embeddings_path, name + '_payload_titles.msgpack') |
|
dense_path = os.path.join(embeddings_path, name + '_dense.npz') |
|
sparse_path = os.path.join(embeddings_path, name + '_sparse.npz') |
|
sparse_titles_path = os.path.join(embeddings_path, name + '_sparse_titles.npz') |
|
ids_path = os.path.join(embeddings_path, name + '_ids.npy') |
|
ids_titles_path = os.path.join(embeddings_path, name + '_ids_titles.npy') |
|
|
|
if not os.path.exists(embeddings_path): |
|
os.mkdir(embeddings_path) |
|
|
|
st.write('Downloading and Loading Video Games Dataset coming from Wikipedia...') |
|
|
|
dataset = pd.read_parquet(os.path.join(os.getenv('HOME'),'data', 'train_pages.parquet.zst'), engine='pyarrow') |
|
|
|
for columnName, columnData in dataset.iteritems(): |
|
if columnName == 'text': |
|
documents = columnData.values.tolist() |
|
elif columnName == 'section_title': |
|
metadatas_titles = columnData.values.tolist() |
|
elif columnName == 'url': |
|
metadatas_url = columnData.values.tolist() |
|
|
|
st.write('Transforming the Wikipedia Video Games Dataset into ingestable format for our Qdrant Vector Database...') |
|
|
|
payload_docs = [{ 'text': text, 'metadata': { 'url': url } } for text, url in zip(documents, metadatas_url)] |
|
|
|
start_sparse = time.time() |
|
|
|
sparse_embeddings = [SparseVector(indices=s.indices.tolist(), values=s.values.tolist()) for s in sparse_model.embed(metadatas_titles, 32)] |
|
|
|
end_sparse = time.time() |
|
final_sparse = end_sparse - start_sparse |
|
print(f'SPARSE TIME: {final_sparse}') |
|
|
|
st.write('Saving on disk the Wikipedia Video Games Dataset into quickly ingestable format...') |
|
|
|
with open(payload_titles_path, "wb") as outfile_texts: |
|
packed_payload = msgpack.packb(payload_docs, use_bin_type=True) |
|
outfile_texts.write(packed_payload) |
|
|
|
max_index = 0 |
|
for embedding in sparse_embeddings: |
|
if len(embedding.indices) > 0: |
|
max_index = max(max_index, max(embedding.indices)) |
|
|
|
sparse_matrices = [] |
|
for embedding in sparse_embeddings: |
|
data = np.array(embedding.values) |
|
indices = np.array(embedding.indices) |
|
indptr = np.array([0, len(data)]) |
|
matrix = csr_matrix((data, indices, indptr), shape=(1, max_index + 1)) |
|
sparse_matrices.append(matrix) |
|
|
|
combined_sparse_matrix = vstack(sparse_matrices) |
|
save_npz(sparse_titles_path, combined_sparse_matrix) |
|
|
|
unique_ids = [] |
|
while len(unique_ids) < len(payload_docs): |
|
new_id = uuid.uuid4() |
|
while new_id.hex[0] == '0': |
|
new_id = uuid.uuid4() |
|
unique_ids.append(new_id.int) |
|
|
|
np.save(ids_titles_path, np.array(unique_ids), allow_pickle=True) |
|
|
|
st.write('Ingesting saved documents on disk into our Qdrant Vector Database...') |
|
|
|
client.upsert( |
|
collection_name, |
|
points=Batch( |
|
ids=unique_ids, |
|
payloads=payload_docs, |
|
vectors={ |
|
'title-sparse': sparse_embeddings, |
|
} |
|
) |
|
) |
|
|
|
client.update_collection( |
|
collection_name=collection_name, |
|
optimizer_config=OptimizersConfigDiff(indexing_threshold=20000) |
|
) |
|
else: |
|
st.write('Loading the saved documents on disk') |
|
|
|
patterns = ['_ids', '_payload', '_dense', '_sparse'] |
|
|
|
unique_ids, payload_docs, dense_embeddings, sparse_embeddings = [ |
|
collect_files(embeddings_path, pattern) for pattern in patterns |
|
] |
|
|
|
patterns = ['_ids_titles', '_payload_titles', '_sparse_titles'] |
|
|
|
unique_ids_titles, payload_docs_titles, sparse_embeddings_titles = [ |
|
collect_files(embeddings_path, pattern) for pattern in patterns |
|
] |
|
|
|
st.write('Ingesting saved documents on disk into our Qdrant Vector Database...') |
|
|
|
print(f'LEN : {len(unique_ids)} {len(payload_docs)} {len(dense_embeddings)} {len(sparse_embeddings)} {len(unique_ids_titles)} {len(payload_docs_titles)} {len(sparse_embeddings_titles)}') |
|
|
|
if unique_ids: |
|
client.upsert( |
|
collection_name, |
|
points=Batch( |
|
ids=unique_ids, |
|
payloads=payload_docs, |
|
vectors={ |
|
'text-dense': dense_embeddings, |
|
'text-sparse': sparse_embeddings |
|
} |
|
) |
|
) |
|
|
|
client.upsert( |
|
collection_name, |
|
points=Batch( |
|
ids=unique_ids_titles, |
|
payloads=payload_docs_titles, |
|
vectors={ |
|
'title-sparse': sparse_embeddings_titles |
|
} |
|
) |
|
) |
|
|
|
client.update_collection( |
|
collection_name=collection_name, |
|
optimizer_config=OptimizersConfigDiff(indexing_threshold=20000) |
|
) |
|
|
|
st.write('Building FSM Index for Agentic Behaviour of our AI...') |
|
|
|
answer = generate_answer( |
|
'aggro', |
|
client, |
|
collection_name, |
|
dense_model, |
|
sparse_model, |
|
'', |
|
'Exact Search', |
|
False, |
|
gen_text, |
|
gen_context_choice, |
|
gen_question_choice |
|
) |
|
|
|
status.update( |
|
label="Processing Complete!", state="complete", expanded=False |
|
) |
|
|
|
time.sleep(5) |
|
container.empty() |
|
|
|
return client, collection_name, dense_model, sparse_model, gen_text, gen_context_choice, gen_question_choice |
|
|
|
|
|
if __name__ == '__main__': |
|
st.set_page_config(page_title="Multipurpose AI Agent",layout="wide", initial_sidebar_state='auto') |
|
|
|
client, collection_name, dense_model, sparse_model, gen_text, gen_context_choice, gen_question_choice = load_models_and_documents() |
|
|
|
styles = { |
|
"nav": { |
|
"background-color": "rgb(204, 200, 194)", |
|
}, |
|
"div": { |
|
"max-width": "32rem", |
|
}, |
|
"span": { |
|
"border-radius": "0.5rem", |
|
"color": "rgb(125, 102, 84)", |
|
"margin": "0 0.125rem", |
|
"padding": "0.4375rem 0.625rem", |
|
}, |
|
"active": { |
|
"background-color": "rgba(255, 255, 255, 0.25)", |
|
}, |
|
"hover": { |
|
"background-color": "rgba(255, 255, 255, 0.35)", |
|
}, |
|
} |
|
|
|
if 'menu_id' not in st.session_state: |
|
st.session_state.menu_id = 'ChatBot' |
|
|
|
st.session_state.menu_id = st_navbar( |
|
['ChatBot', 'Documents'], |
|
st.session_state.menu_id, |
|
options={ |
|
'hide_nav': False, |
|
'fix_shadow': False, |
|
'use_padding': False |
|
}, |
|
styles=styles |
|
) |
|
|
|
st.title('Multipurpose AI Agent') |
|
|
|
|
|
data_editor_path = os.path.join(os.getenv('HF_HOME'), 'documents') |
|
|
|
if 'df' not in st.session_state: |
|
if os.path.exists(data_editor_path): |
|
st.session_state.df = pd.read_parquet(os.path.join(data_editor_path, 'data_editor.parquet.lz4'), engine='pyarrow') |
|
else: |
|
st.session_state.df = pd.DataFrame() |
|
os.mkdir(data_editor_path) |
|
st.session_state.df.to_parquet( |
|
os.path.join( |
|
data_editor_path, |
|
'data_editor.parquet.lz4' |
|
), |
|
compression='lz4', |
|
engine='pyarrow' |
|
) |
|
|
|
def on_change_data_editor(client, collection_name): |
|
print(f'Check : {st.session_state.key_data_editor}') |
|
|
|
if st.session_state.key_data_editor['deleted_rows']: |
|
st.toast('Wait for deletion to complete...') |
|
embeddings_path = os.path.join(os.getenv('HF_HOME'), 'embeddings') |
|
|
|
for deleted_idx in st.session_state.key_data_editor['deleted_rows']: |
|
name = st.session_state.df.loc[deleted_idx, 'document'] |
|
print(f'WHAT IS THAT : {name}') |
|
|
|
os.remove(os.path.join(embeddings_path, name + '_ids.npy')) |
|
os.remove(os.path.join(embeddings_path, name + '_payload.msgpack')) |
|
os.remove(os.path.join(embeddings_path, name + '_dense.npz')) |
|
os.remove(os.path.join(embeddings_path, name + '_sparse.npz')) |
|
|
|
client.delete( |
|
collection_name=collection_name, |
|
points_selector=Filter( |
|
must=[ |
|
FieldCondition( |
|
key='url', |
|
match=MatchValue(value=st.session_state.df.loc[deleted_idx, 'path']) |
|
) |
|
] |
|
) |
|
) |
|
|
|
st.session_state.df.drop(deleted_idx) |
|
|
|
st.toast('Deletion Completed !', icon='🎉') |
|
|
|
if st.session_state.menu_id == 'Documents': |
|
with st.sidebar: |
|
if "cached_files" not in st.session_state: |
|
st.session_state.cached_files = [] |
|
|
|
uploaded_files = st.file_uploader("Upload a file :", accept_multiple_files=True, type=['pptx', 'ppt']) |
|
|
|
for uploaded_file in uploaded_files: |
|
|
|
if uploaded_file not in st.session_state.cached_files: |
|
st.session_state.cached_files.append(uploaded_file) |
|
|
|
file_name = os.path.basename(uploaded_file.name) |
|
base_name, ext = os.path.splitext(file_name) |
|
|
|
processing_time = datetime.now().strftime('%d %b %Y, %I:%M %p') |
|
|
|
full_path = os.path.realpath(uploaded_file.name) |
|
file_type = ext.lstrip('.') |
|
|
|
d = { |
|
'icon': icon_to_types[file_type][0], |
|
'document': base_name, |
|
'type': icon_to_types[file_type][1], |
|
'path': full_path, |
|
'time': [datetime.strptime(processing_time, '%d %b %Y, %I:%M %p')], |
|
'exact_search': False |
|
} |
|
|
|
if (st.session_state.df.empty) or (base_name not in st.session_state.df['document'].tolist()): |
|
st.session_state.df = pd.concat( |
|
[ |
|
st.session_state.df, |
|
pd.DataFrame(data={ |
|
'icon': icon_to_types[file_type][0], |
|
'document': base_name, |
|
'type': icon_to_types[file_type][1], |
|
'path': full_path, |
|
'time': [datetime.strptime(processing_time, '%d %b %Y, %I:%M %p')], |
|
'exact_search': False |
|
}) |
|
], |
|
ignore_index=True |
|
) |
|
else: |
|
idx = st.session_state.df.index[st.session_state.df['document']==base_name].tolist()[0] |
|
st.session_state.df.loc[idx] = { |
|
'icon': icon_to_types[file_type][0], |
|
'document': base_name, |
|
'type': icon_to_types[file_type][1], |
|
'path': full_path, |
|
'time': datetime.strptime(processing_time, '%d %b %Y, %I:%M %p'), |
|
'exact_search': False |
|
} |
|
|
|
st.session_state.df.to_parquet( |
|
os.path.join( |
|
data_editor_path, |
|
'data_editor.parquet.lz4' |
|
), |
|
compression='lz4', |
|
engine='pyarrow' |
|
) |
|
|
|
documents, ids = ppt_chunker(uploaded_file) |
|
|
|
dense_time = time.time() |
|
dense_embeddings, tokens_count = asyncio.run(embed_text(dense_model[0], documents)) |
|
end_dense_time = time.time() |
|
final_dense = end_dense_time - dense_time |
|
|
|
print(f'DENSE TIME : {final_dense}') |
|
sparse_embeddings = [s for s in sparse_model.embed(documents, 32)] |
|
|
|
embeddings_path = os.path.join(os.getenv('HF_HOME'), 'embeddings') |
|
|
|
def generate_unique_id(existing_ids): |
|
while True: |
|
new_id = uuid.uuid4() |
|
while new_id.hex[0] == '0': |
|
new_id = uuid.uuid4() |
|
new_id = new_id.int |
|
if new_id not in existing_ids: |
|
return new_id |
|
|
|
for filename in os.listdir(embeddings_path): |
|
if '_ids' in filename: |
|
list_ids = np.load(os.path.join(embeddings_path, filename), allow_pickle=True).tolist() |
|
|
|
for i, ids_ in enumerate(ids): |
|
if ids_ in list_ids: |
|
ids[i] = generate_unique_id(list_ids) |
|
|
|
metadatas_list = [{'url': full_path}] * len(documents) |
|
|
|
payload_docs = [{ 'text': documents[i], 'metadata': metadata } for i, metadata in enumerate(metadatas_list)] |
|
|
|
client.upsert( |
|
collection_name=collection_name, |
|
points=Batch( |
|
ids=ids, |
|
payloads=payload_docs, |
|
vectors={ |
|
'text-dense': dense_embeddings, |
|
'text-sparse': [SparseVector(indices=s.indices.tolist(), values=s.values.tolist()) for s in sparse_embeddings] |
|
} |
|
) |
|
) |
|
|
|
payload_path = os.path.join(embeddings_path, base_name + '_payload.msgpack') |
|
dense_path = os.path.join(embeddings_path, base_name + '_dense.npz') |
|
sparse_path = os.path.join(embeddings_path, base_name + '_sparse.npz') |
|
ids_path = os.path.join(embeddings_path, base_name + '_ids.npy') |
|
|
|
with open(payload_path, "wb") as outfile_texts: |
|
packed_payload = msgpack.packb(payload_docs, use_bin_type=True) |
|
outfile_texts.write(packed_payload) |
|
|
|
np.savez_compressed(dense_path, *dense_embeddings) |
|
max_index = 0 |
|
for embedding in sparse_embeddings: |
|
if len(embedding.indices) > 0: |
|
max_index = max(max_index, max(embedding.indices)) |
|
|
|
sparse_matrices = [] |
|
|
|
for embedding in sparse_embeddings: |
|
data = np.array(embedding.values) |
|
indices = np.array(embedding.indices) |
|
indptr = np.array([0, len(data)]) |
|
matrix = csr_matrix((data, indices, indptr), shape=(1, max_index + 1)) |
|
sparse_matrices.append(matrix) |
|
|
|
combined_sparse_matrix = vstack(sparse_matrices) |
|
save_npz(sparse_path, combined_sparse_matrix) |
|
|
|
np.save(ids_path, np.array(ids), allow_pickle=True) |
|
|
|
st.toast('Document(s) Ingested !', icon='🎉') |
|
st.rerun() |
|
|
|
st.session_state.df = st.data_editor( |
|
st.session_state.df, |
|
num_rows="dynamic", |
|
use_container_width=True, |
|
hide_index=True, |
|
on_change=on_change_data_editor, |
|
args=(client, collection_name), |
|
key='key_data_editor', |
|
column_config={ |
|
'icon': st.column_config.ImageColumn( |
|
'Document' |
|
), |
|
"document": st.column_config.TextColumn( |
|
"Name", |
|
help="Name of the document", |
|
required=True, |
|
disabled=True |
|
), |
|
"type": st.column_config.SelectboxColumn( |
|
'File type', |
|
help='The file format extension of this document', |
|
required=True, |
|
options=[ |
|
'Powerpoint', |
|
'Microsoft Word', |
|
'Excel' |
|
], |
|
disabled=True |
|
), |
|
"path": st.column_config.TextColumn( |
|
'Path', |
|
help='Path to the document', |
|
required=False, |
|
disabled=True |
|
), |
|
"time": st.column_config.DatetimeColumn( |
|
'Date and hour', |
|
help='When this document has been ingested here for the last time', |
|
format="D MMM YYYY, h:mm a", |
|
required=True, |
|
disabled=True |
|
), |
|
"exact_search": st.column_config.CheckboxColumn( |
|
'Exact Search', |
|
help='Wether the Exact Search is available for this document or not', |
|
required=True, |
|
default=False, |
|
disabled=True |
|
) |
|
} |
|
) |
|
|
|
conversations_path = os.path.join(os.getenv('HF_HOME'), 'conversations') |
|
try: |
|
with open(conversations_path, 'rb') as fp: |
|
packed_bytes = fp.read() |
|
conversations: Dict[str, list] = msgpack.unpackb(packed_bytes, raw=False) |
|
except: |
|
conversations = {} |
|
|
|
if st.session_state.menu_id == 'ChatBot': |
|
if 'id_chat' not in st.session_state: |
|
st.session_state.id_chat = 'New Conversation' |
|
if 'search_strategy' not in st.session_state: |
|
st.session_state.search_strategy = None |
|
if 'local_user_input' not in st.session_state: |
|
st.session_state.local_user_input = None |
|
|
|
def on_change_documents_only(): |
|
if st.session_state.documents_only: |
|
st.session_state.toggle_docs = { |
|
'tooltip': 'The AI answer your questions only considering the documents provided', |
|
'display': True |
|
} |
|
else: |
|
st.session_state.toggle_docs = { |
|
'tooltip': """The AI answer your questions considering the documents provided, and if it doesn't found the answer in them, try to find in its own internal knowledge""", |
|
'display': False |
|
} |
|
|
|
def options_list(conversations: Dict[str, list]): |
|
if st.session_state.id_chat == 'New Conversation': |
|
return [st.session_state.id_chat] + list(conversations.keys()) |
|
else: |
|
return ['New Conversation'] + list(conversations.keys()) |
|
|
|
with st.sidebar: |
|
if 'toggle_docs' not in st.session_state: |
|
st.session_state.toggle_docs = { |
|
'tooltip': 'The AI answer your questions only considering the documents provided', |
|
'display': True |
|
} |
|
|
|
st.toggle( |
|
label="""Enable 'Documents-Only' Mode""", |
|
value=st.session_state.toggle_docs['display'], |
|
on_change=on_change_documents_only, |
|
key="documents_only", |
|
help=st.session_state.toggle_docs['tooltip'] |
|
) |
|
|
|
st.divider() |
|
|
|
st.session_state.id_chat = st.selectbox( |
|
label='Choose a Conversation', |
|
options=options_list(conversations), |
|
index=0, |
|
key='chat_id' |
|
) |
|
|
|
st.session_state.messages = conversations[st.session_state.id_chat] if st.session_state.id_chat != 'New Conversation' else [] |
|
|
|
def update_selectbox_remove(conversations_path, conversations): |
|
conversations.pop(st.session_state.chat_id) |
|
with open(conversations_path, 'wb') as fp: |
|
packed_bytes = msgpack.packb(conversations, use_bin_type=True) |
|
fp.write(packed_bytes) |
|
st.session_state.chat_id = 'New Conversation' |
|
|
|
st.button( |
|
'Delete Conversation', |
|
use_container_width=True, |
|
disabled=False if st.session_state.id_chat != 'New Conversation' else True, |
|
on_click=update_selectbox_remove, |
|
args=(conversations_path, conversations) |
|
) |
|
|
|
for message in st.session_state.messages: |
|
with st.chat_message(message["role"]): |
|
st.markdown(message["content"]) |
|
|
|
def generate_conv_title(generator): |
|
st.session_state.local_user_input = st.session_state.user_input |
|
print(f'USER INPUT : {st.session_state.user_input}') |
|
st.session_state.user_input = " " |
|
if st.session_state.chat_id == 'New Conversation': |
|
output = json.loads( |
|
generator( |
|
build_prompt_conv(st.session_state.local_user_input.replace(st.session_state.search_strategy, "") if st.session_state.search_strategy != None else st.session_state.local_user_input), |
|
max_tokens=10, |
|
sampling_params=SamplingParams(temperature=0, top_p=0.9, max_tokens=10, top_k=10) |
|
) |
|
) |
|
print(f'OUTPUT : {output}') |
|
st.session_state.chat_id = output |
|
st.session_state.messages = [] |
|
|
|
torch.cuda.empty_cache() |
|
|
|
conversations.update({st.session_state.chat_id: st.session_state.messages}) |
|
with open(conversations_path, 'wb') as fp: |
|
packed_bytes = msgpack.packb(conversations, use_bin_type=True) |
|
fp.write(packed_bytes) |
|
|
|
with _bottom.container(): |
|
st.session_state.search_strategy = pills("", ["Exact Search", "Explain Further"], ["🎯", "🚀"], index=None) |
|
|
|
if prompt := _bottom.text_input( |
|
" ", |
|
value=st.session_state.search_strategy + ' : ' if st.session_state.search_strategy != None else "", |
|
on_change=generate_conv_title, |
|
key='user_input', |
|
placeholder='Message Video Game Assistant', |
|
label_visibility='collapsed', |
|
args=(gen_text, ) |
|
): |
|
if prompt != ('Exact Search : ' or 'Explain Further : '): |
|
st.chat_message("user").markdown(st.session_state.local_user_input) |
|
st.session_state.messages.append({"role": "user", "content": st.session_state.local_user_input}) |
|
|
|
ai_response = generate_answer( |
|
st.session_state.local_user_input, |
|
client, |
|
collection_name, |
|
dense_model, |
|
sparse_model, |
|
"\n".join([f'{msg["role"]}: {msg["content"]}' for msg in st.session_state.messages]), |
|
st.session_state.search_strategy, |
|
st.session_state.documents_only, |
|
gen_text, |
|
gen_context_choice, |
|
gen_question_choice |
|
) |
|
|
|
with st.chat_message("assistant"): |
|
message_placeholder = st.empty() |
|
full_response = "" |
|
for chunk in re.split(r'(\s+)', ai_response): |
|
full_response += chunk + " " |
|
time.sleep(0.05) |
|
message_placeholder.write(full_response + '▌') |
|
message_placeholder.write(re.sub('▌', '', full_response)) |
|
|
|
st.session_state.messages.append({"role": "assistant", "content": full_response}) |
|
|
|
conversations.update({st.session_state.id_chat: st.session_state.messages}) |
|
with open(conversations_path, 'wb') as fp: |
|
packed_bytes = msgpack.packb(conversations, use_bin_type=True) |
|
fp.write(packed_bytes) |
