Update app.py

app.py

@@ -13,11 +13,9 @@ import numpy as np
 import pandas as pd
 import streamlit as st
 
-from vllm import LLM
 from datetime import datetime
 from streamlit import _bottom
 from pydantic import BaseModel
-from outlines.models import VLLM
 from streamlit_pills import pills
 from dense_embed import embed_text
 from ppt_chunker import ppt_chunker
@@ -103,12 +101,14 @@ icon_to_types = {
 def generate_answer(query: str, 
                     client: QdrantClient, 
                     collection_name: str, 
-                    llm, 
                     dense_model: AsyncEmbeddingEngine, 
                     sparse_model: SparseTextEmbedding, 
                     past_messages: str, 
                     search_strategy,
-                    documents_only: bool
+                    documents_only: bool,
+                    gen_text,
+                    gen_context_choice,
+                    gen_question_choice
                    ):
     sparse_embeddings = list(sparse_model.query_embed(query))[0].as_object()
     
@@ -130,13 +130,9 @@ def generate_answer(query: str,
         contents, metadatas = [list(t) for t in zip(*docs)]
             
         context = "\n".join(contents)
-        
-        regex = build_regex_from_schema(schema, r"[\n ]?")
-        gen_text = outlines.generate.regex(llm, regex)
 
-        gen_choice = outlines.generate.choice(llm, choices=['Yes', 'No'])
         prompt = route_llm(context, query)
-        action = gen_choice(prompt, max_tokens=2, sampling_params=SamplingParams(temperature=0))
+        action = gen_context_choice(prompt, max_tokens=2, sampling_params=SamplingParams(temperature=0))
         print(f'Choice: {action}')
     
         if action == 'Yes':      
@@ -155,9 +151,8 @@ def generate_answer(query: str,
             if documents_only == False:
                 answer = f'Documents Based :\n\n{answer}'
         else:
-            gen_choice = outlines.generate.choice(llm, choices=['Domain-Specific Question', 'General Question'])
             prompt = question_type_prompt(query)
-            action = gen_choice(prompt, max_tokens=3, sampling_params=SamplingParams(temperature=0))
+            action = gen_question_choice(prompt, max_tokens=3, sampling_params=SamplingParams(temperature=0))
             print(f'Choice 2: {action}')
             
             if action == 'General Question':
@@ -210,6 +205,7 @@ def collect_files(directory, pattern):
                 
     return array
 
+@st.cache_resource(show_spinner=False)
 def load_models_and_documents():
     container = st.empty()
     
@@ -240,17 +236,22 @@ def load_models_and_documents():
 
         st.write('Downloading and Loading Mistral Nemo quantized with GPTQ and using Outlines + vLLM Engine as backend...')
 
-        llm = LLM(
-            model="shuyuej/Mistral-Nemo-Instruct-2407-GPTQ",
+        llm = outlines.models.vllm(
+            model_name="shuyuej/Mistral-Nemo-Instruct-2407-GPTQ",
             tensor_parallel_size=1,
             enforce_eager=True,
-            gpu_memory_utilization=0.9,
+            gpu_memory_utilization=1,
             max_model_len=8192,
             dtype=torch.float16,
             max_num_seqs=64,
             quantization="gptq"
         )
-        model = VLLM(llm)
+
+        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...')
 
@@ -441,7 +442,19 @@ def load_models_and_documents():
 
         st.write('Building FSM Index for Agentic Behaviour of our AI...')
 
-        answer = generate_answer('aggro', client, collection_name, model, dense_model, sparse_model, '', 'Exact Search', False)
+        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
@@ -450,14 +463,13 @@ def load_models_and_documents():
     time.sleep(5)
     container.empty()
     
-    return client, collection_name, llm, model, dense_model, sparse_model
+    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')
 
-    if 'client' not in st.session_state:
-        st.session_state.client, st.session_state.collection_name, st.session_state.llm, st.session_state.model, st.session_state.dense_model, st.session_state.sparse_model = load_models_and_documents()
+    client, collection_name, dense_model, sparse_model, gen_text, gen_context_choice, gen_question_choice = load_models_and_documents()
 
     styles = {
         "nav": {
@@ -530,7 +542,7 @@ if __name__ == '__main__':
                 os.remove(os.path.join(embeddings_path, name + '_dense.npz'))
                 os.remove(os.path.join(embeddings_path, name + '_sparse.npz'))
                 
-                st.session_state.client.delete(
+                client.delete(
                     collection_name=collection_name,
                     points_selector=Filter(
                         must=[
@@ -610,10 +622,10 @@ if __name__ == '__main__':
                         engine='pyarrow'
                     )
                     
-                    documents, ids = ppt_chunker(uploaded_file, st.session_state.llm)
+                    documents, ids = ppt_chunker(uploaded_file, llm)
 
-                    dense_embeddings, tokens_count = asyncio.run(embed_text(st.session_state.dense_model[0], documents))
-                    sparse_embeddings = [s for s in st.session_state.sparse_model.embed(documents, 32)]
+                    dense_embeddings, tokens_count = asyncio.run(embed_text(dense_model[0], documents))
+                    sparse_embeddings = [s for s in sparse_model.embed(documents, 32)]
 
                     embeddings_path = os.path.join(os.getenv('HF_HOME'), 'embeddings')
 
@@ -638,8 +650,8 @@ if __name__ == '__main__':
 
                     payload_docs = [{ 'text': documents[i], 'metadata': metadata } for i, metadata in enumerate(metadatas_list)]
                     
-                    st.session_state.client.upsert(
-                        collection_name=st.session_state.collection_name,
+                    client.upsert(
+                        collection_name=collection_name,
                         points=Batch(
                             ids=ids,
                             payloads=payload_docs,
@@ -687,7 +699,7 @@ if __name__ == '__main__':
             use_container_width=True,
             hide_index=True,
             on_change=on_change_data_editor,
-            args=(st.session_state.client, st.session_state.collection_name),
+            args=(client, collection_name),
             key='key_data_editor',
             column_config={
                 'icon': st.column_config.ImageColumn(
@@ -812,17 +824,20 @@ if __name__ == '__main__':
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
 
-        def generate_conv_title(llm):
+        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 = llm.chat(
-                    build_prompt_conv(st.session_state.local_user_input),
-                    SamplingParams(temperature=0, top_p=0.9, max_tokens=10, top_k=10)
+                output = json.loads(
+                    generator(
+                        build_prompt_conv(st.session_state.local_user_input),
+                        max_tokens=10,
+                        sampling_params=SamplingParams(temperature=0, top_p=0.9, top_k=10)
+                    )
                 )
-                print(f'OUTPUT : {output[0].outputs[0].text}')
-                st.session_state.chat_id = output[0].outputs[0].text.replace('"', '')
+                print(f'OUTPUT : {output}')
+                st.session_state.chat_id = output
                 st.session_state.messages = []
 
                 torch.cuda.empty_cache()
@@ -842,7 +857,7 @@ if __name__ == '__main__':
             key='user_input',
             placeholder='Message Video Game Assistant',
             label_visibility='collapsed',
-            args=(st.session_state.llm, )
+            args=(gen_text, )
         ):
             if prompt != ('Exact Search : ' or 'Explain Further : '):
                 st.chat_message("user").markdown(st.session_state.local_user_input)
@@ -850,14 +865,16 @@ if __name__ == '__main__':
                 
                 ai_response = generate_answer(
                     st.session_state.local_user_input, 
-                    st.session_state.client, 
-                    st.session_state.collection_name, 
-                    st.session_state.model, 
-                    st.session_state.dense_model, 
-                    st.session_state.sparse_model, 
+                    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
+                    st.session_state.documents_only,
+                    gen_text,
+                    gen_context_choice,
+                    gen_question_choice
                 )
                 
                 with st.chat_message("assistant"):