import os import re import lz4 import nltk import copy import time import vllm import torch import spacy import shutil import msgpack import tempfile import threading import validators import numpy as np import pandas as pd import streamlit as st from pathlib import Path from numpy import ndarray from outlines import models from datetime import datetime from typing import List, Dict from transformers import AutoTokenizer from qdrant_client import QdrantClient from optimum_encoder import OptimumEncoder from streamlit_navigation_bar import st_navbar from ppt_chunker import ppt_chunk from unstructured.cleaners.core import clean from unstructured.partition.pptx import partition_pptx from fastembed import SparseTextEmbedding, SparseEmbedding from unstructured.nlp.tokenize import download_nltk_packages from huggingface_hub import snapshot_download, hf_hub_download from scipy.sparse import csr_matrix, save_npz, load_npz, vstack from langchain_experimental.text_splitter import SemanticChunker from langchain_community.document_loaders import WikipediaLoader, WebBaseLoader from qdrant_client.models import ( NamedSparseVector, NamedVector, SparseVector, PointStruct, ScoredPoint, Prefetch, FusionQuery, Fusion, SearchRequest, Modifier, OptimizersConfigDiff, HnswConfigDiff, Distance, VectorParams, SparseVectorParams, SparseIndexParams, Batch ) 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 transform_query(query: str) -> str: """ For retrieval, add the prompt for query (not for documents). """ return f'Represent this sentence for searching relevant passages: {query}' def query_hybrid_search(query: str, client: QdrantClient, collection_name: str, dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding): dense_embeddings = dense_model.embed_query(transform_query(query))[0] sparse_embeddings = list(sparse_model.query_embed(query))[0] return client.query_points( collection_name=collection_name, prefetch=[ Prefetch(query=sparse_embeddings, using="text-sparse", limit=10), Prefetch(query=dense_embeddings, using="text-dense", limit=10) ], query=FusionQuery(fusion=Fusion.RRF), limit=3 ) def main(query: str, client: QdrantClient, collection_name: str, tokenizer: AutoTokenizer, llm: vllm.LLM, dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding): scored_points = query_hybrid_search(query, client, collection_name, dense_model, sparse_model).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) seen_values = set() result_metadatas = "\n\n".join( f'{value}' for metadata in metadatas for key, value in metadata.items() if (value not in seen_values and not seen_values.add(value)) ) args = {'context': context, 'query': query} messages = [ {"role": "system", "content": 'You are a helpful assistant.'}, {"role": "user", "content": st.session_state.toggle_docs['qa_prompt'].format(**args)} ] prompts = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) outputs = llm.generate( prompts=prompts, sampling_params=vllm.SamplingParams( temperature=0, max_tokens=3000 ) ) print(f'TEXT: {outputs}') text = outputs[0].outputs[0].text messages_2 = [ {"role": "system", "content": """Act like a professional summary writer. You have been providing summarization services for various types of documents, including academic papers, legal texts, and business reports, for over 20 years. Your expertise includes extracting key points and important details concisely without adding unnecessary introductory phrases.""" }, {"role": "user", "content": f"""Write a summary of the following text delimited by triple backquotes. Ensure the summary covers the key points of the text. Do not introduce the summary with sentences like "Here is the summary:" or similar. The summary should be detailed, precise, and directly convey the essential information from the text. ```{text}``` Let's think step-by-step.""" } ] prompts_2 = tokenizer.apply_chat_template(messages_2, tokenize=False, add_generation_prompt=True) outputs_2 = llm.generate( prompts=prompts_2, sampling_params=vllm.SamplingParams( temperature=0.3, max_tokens=3000 ) ) answer = outputs_2[0].outputs[0].text answer_with_metadatas = f"{answer}\n\n\nSource(s) :\n\n{result_metadatas}" if st.session_state.documents_only: return answer if 'no_answer' in text else answer_with_metadatas else: return f'Internal Knowledge :\n\n{answer}' if 'knowledge_topic' in text else f'Documents Based :\n\n{answer_with_metadatas}' @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 Mistral Nemo AI Model...') model_path = snapshot_download('casperhansen/mistral-nemo-instruct-2407-awq') st.write('Loading Mistral Nemo AI Model with vLLM backend...') tokenizer = AutoTokenizer.from_pretrained(model_path) llm = vllm.LLM( model=model_path, tensor_parallel_size=1, trust_remote_code=True, enforce_eager=True, quantization="awq", gpu_memory_utilization=0.9, max_model_len=11000, dtype=torch.float16 ) model = models.VLLM(llm) st.write('Downloading and Loading Embeddings Models...') dense_model = OptimumEncoder( device="cuda", cache_dir=os.getenv('HF_HOME') ) sparse_model = SparseTextEmbedding( 'Qdrant/bm42-all-minilm-l6-v2-attentions', cache_dir=os.getenv('HF_HOME'), providers=['CPUExecutionProvider'] ) st.write('Loading Natural Language Processing Model for English...') nlp = spacy.load("en_core_web_sm") download_nltk_packages() st.write('Creating Collection for our Qdrant Vector Database in RAM memory...') client = QdrantClient(':memory:') collection_name = 'collection_demo' client.create_collection( collection_name, { "text-dense": VectorParams( size=1024, distance=Distance.COSINE, on_disk=False ) }, { "text-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 ) ) name = 'action_rpg' embeddings_path = os.path.join(os.getenv('HF_HOME'), 'embeddings') payload_path = os.path.join(embeddings_path, name + '_payload') dense_path = os.path.join(embeddings_path, name + '_dense.npz') sparse_path = os.path.join(embeddings_path, name + '_sparse.npz') if not os.path.exists(embeddings_path): os.mkdir(embeddings_path) st.write('Downloading and Loading Video Games Dataset coming from Wikipedia...') docs_1 = WikipediaLoader(query='Action-RPG').load() docs_2 = WikipediaLoader(query='Real-time strategy').load() docs_3 = WikipediaLoader(query='First-person shooter').load() docs_4 = WikipediaLoader(query='Multiplayer online battle arena').load() docs_5 = WikipediaLoader(query='List of video game genres').load() docs = docs_1 + docs_2 + docs_3 + docs_4 + docs_5 texts, metadatas = [], [] for doc in docs: texts.append(doc.page_content) del doc.metadata['title'] del doc.metadata['summary'] metadatas.append(doc.metadata) st.write('Transforming the Wikipedia Video Games Dataset into ingestable format for our Qdrant Vector Database...') payload_docs, dense_embeddings, sparse_embeddings = chunk_documents(texts, metadatas, dense_model, sparse_model) st.write('Saving on disk the Wikipedia Video Games Dataset into quickly ingestable format...') 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 embedding.indices.size > 0: max_index = max(max_index, np.max(embedding.indices)) sparse_matrices = [] for embedding in sparse_embeddings: data = embedding.values indices = 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) else: st.write('Loading the saved documents on disk') with open(payload_path, "rb") as data_file_payload: decompressed_payload = data_file_payload.read() payload_docs = msgpack.unpackb(decompressed_payload, raw=False) dense_embeddings = list(np.load(dense_path).values()) sparse_embeddings = [] loaded_sparse_matrix = load_npz(sparse_path) for i in range(loaded_sparse_matrix.shape[0]): row = loaded_sparse_matrix.getrow(i) values = row.data indices = row.indices embedding = SparseEmbedding(indices=indices, values=values) sparse_embeddings.append(embedding) st.write('Ingesting saved documents on disk into our Qdrant Vector Database...') client.upsert( collection_name, points=Batch( payload=payload_docs, vectors={ 'text_dense': dense_embeddings, 'text_sparse': [SparseVector(indices=s.indices.tolist(), values=s.values.tolist()) for s in sparse_embeddings] } ) ) client.update_collection( collection_name=collection_name, optimizer_config=OptimizersConfigDiff(indexing_threshold=20000) ) status.update( label="Processing Complete!", state="complete", expanded=False ) time.sleep(5) container.empty() return client, collection_name, tokenizer, model, llm, dense_model, sparse_model, nlp def chunk_documents(texts: List[str], metadatas: List[dict], dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding): text_splitter = SemanticChunker( dense_model, breakpoint_threshold_type='standard_deviation' ) docs = text_splitter.create_documents(texts, metadatas) payload_docs, documents = [], [] for doc in docs: payload_docs.append({ 'text': doc.page_content, 'metadata': doc.metadata }) documents.append(doc.page_content) start_dense = time.time() dense_embeddings = dense_model.embed_documents(documents) end_dense = time.time() final_dense = end_dense - start_dense print(f'DENSE TIME: {final_dense}') start_sparse = time.time() sparse_embeddings = [s for s in sparse_model.embed(documents, 32)] end_sparse = time.time() final_sparse = end_sparse - start_sparse print(f'SPARSE TIME: {final_sparse}') return payload_docs, dense_embeddings, sparse_embeddings def on_change_documents_only(): if st.session_state.documents_only: st.session_state.toggle_docs = { 'qa_prompt': """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, reply with 'no_answer'. Use three sentences maximum and keep the answer concise. Question: {query} Context: {context} Answer:""", 'tooltip': 'The AI answer your questions only considering the documents provided', 'display': True } else: st.session_state.toggle_docs = { 'qa_prompt': """If the context is not relevant, please answer the question by using your own knowledge about the topic. If you decide to provide information using your own knowledge or general knowledge, write 'knowledge_topic' at the top of your answer. {context} Question: {query}""", '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 } if __name__ == '__main__': st.set_page_config(page_title="Multipurpose AI Agent",layout="wide", initial_sidebar_state='collapsed') client, collection_name, tokenizer, model, llm, dense_model, sparse_model, nlp = load_models_and_documents() if 'menu_id' not in st.session_state: st.session_state.menu_id = 'ChatBot' 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)", }, } 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') #st.markdown("

Multipurpose AI Agent

", unsafe_allow_html=True) if 'df' not in st.session_state: st.session_state.df = pd.DataFrame([0]) if st.session_state.menu_id == 'Documents': st.session_state.df = st.data_editor( st.session_state.df, num_rows="dynamic", use_container_width=True, hide_index=True, column_config={ 'icon': st.column_config.ImageColumn( 'Document' ), "document": st.column_config.TextColumn( "Name", help="Name of the document", required=True ), "type": st.column_config.SelectboxColumn( 'File type', help='The file format extension of this document', required=True, options=[ 'Powerpoint', 'Microsoft Word', 'Excel' ] ), "path": st.column_config.TextColumn( 'Path', help='Path to the document', required=False ), "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 ), "toggle": st.column_config.CheckboxColumn( 'Enable/Disable', help='Either to enable or disable the ability for the ai to find this document', required=True, default=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 Chat' def options_list(conversations: Dict[str, list]): if st.session_state.id_chat == 'New Chat': return [st.session_state.id_chat] + list(conversations.keys()) else: return ['New Chat'] + list(conversations.keys()) with st.sidebar: st.session_state.id_chat = st.selectbox( label='Choose a conversation', options=options_list(conversations), index=0, placeholder='_', key='chat_id' ) st.session_state.messages = conversations[st.session_state.id_chat] if st.session_state.id_chat != 'New Chat' 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 Chat' st.button( 'Delete Chat', use_container_width=True, disabled=False if st.session_state.id_chat != 'New Chat' else True, on_click=update_selectbox_remove, args=(conversations_path, conversations) ) def generate_conv_title(llm: vllm.LLM, tokenizer: AutoTokenizer): if st.session_state.chat_id == 'New Chat': messages = [ {"role": "system", "content": 'You are a helpful assistant.'}, {"role": "user", "content": f"""Understand the question of the user. Resume in one single sentence what is the subject of the conversation and what is the user talking about. Question : {st.session_state.user_input}""" } ] prompts = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) outputs = llm.generate( prompts=prompts, sampling_params=vllm.SamplingParams( temperature=0.3, max_tokens=30 ) ) st.session_state.chat_id = outputs[0].outputs[0].text st.session_state.messages = [] 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) for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) if prompt := st.chat_input( "Message Video Game Assistant", on_submit=generate_conv_title, key='user_input', args=(llm, tokenizer) ): st.chat_message("user").markdown(prompt) st.session_state.messages.append({"role": "user", "content": prompt}) ai_response = main(prompt, client, collection_name, tokenizer, llm, dense_model, sparse_model) 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.2) 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) @st.dialog('URL Ingested !') def url_ingested(): main_thread = threading.main_thread() while True: L = threading.enumerate() L.remove(main_thread) for t in L: t.join() @st.dialog('Documents Ingested !') def documents_ingested(): main_thread = threading.main_thread() while True: L = threading.enumerate() L.remove(main_thread) for t in L: t.join() with st.sidebar: st.divider() if 'toggle_docs' not in st.session_state: st.session_state.toggle_docs = { 'qa_prompt': """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, reply with 'no_answer'. Use three sentences maximum and keep the answer concise. Question: {query} Context: {context} Answer:""", '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() url = st.text_input("Scrape an URL link :") if validators.url(url): docs = WebBaseLoader(url).load() print(f'WebBaseLoader: {docs[0].metadata}') texts, metadatas = [], [] for doc in docs: texts.append(doc.page_content) del doc.metadata['title'] del doc.metadata['description'] del doc.metadata['language'] metadatas.append(doc.metadata) payload_docs, dense_embeddings, sparse_embeddings = chunk_documents(texts, metadatas, dense_model, sparse_model) client.upsert( collection_name, make_points( texts, metadatas, dense_embeddings, sparse_embeddings ) ) url_ingested() st.divider() uploaded_files = st.file_uploader("Upload a file :", accept_multiple_files=True, type=['pptx', 'ppt']) print(f'uploaded-files : {uploaded_files}') for uploaded_file in uploaded_files: processing_time = datetime.now().strftime('%d %b %Y, %I:%M %p') file_name = os.path.basename(uploaded_file.name) base_name, ext = os.path.splitext(file_name) 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')], 'toggle': True } st.session_state.df = pd.concat([st.session_state.df, pd.DataFrame(data=d)]) elements = partition_pptx(file=uploaded_file) for elem in elements: elem.text = clean(elem.text, bullets=True) text_type = elem.to_dict()['type'] print(f'UNSTRUCTURED TEXT: {text_type} , {elem.text}') weakDict, tables = ppt_chunk(uploaded_file, nlp) documents = weakDict.all_texts() dense = dense_model.embed_documents(documents) print(f'dense: {type(dense)}') sparse = [SparseVector(indices=s.indices.tolist(), values=s.values.tolist()) for s in sparse_model.embed(documents, 32)] print(f'dense: {type(sparse)}') client.upsert( collection_name=collection_name, points=Batch( ids=weakDict.all_ids(), payload= payload={ 'texts': documents, 'metadatas': weakDict.all_metadatas() }, vectors={ 'text-dense': dense, 'text-sparse': sparse } ) ) #client.upsert( # collection_name, # make_points( # documents, # weakDict.all_metadatas(), # dense_embeddings, # sparse_embeddings # ) #) documents_ingested()