Update app.py

app.py

@@ -16,6 +16,7 @@ import numpy as np
 import pandas as pd
 import streamlit as st
 from vllm import LLM
+from json import pyjson
 from numpy import ndarray
 from datetime import datetime
 from typing import List, Dict
@@ -72,7 +73,8 @@ from qdrant_client.models import (
 global_state_documents_only = False
 
 class Question(BaseModel):
-    answer: str = Field(..., description='Answer to the Query')
+    answer: str
+schema = pyjson.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==',
@@ -224,8 +226,9 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
         
     context = "\n".join(contents)
     print(f'Context : \n + {context}')
-    
-    gen_text = outlines.generate.json(llm, Question, whitespace_pattern=r"[\n ]?")
+
+    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)
@@ -242,9 +245,8 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
         result_metadatas = "\n\n".join(f'{value}' for value in filtered_metadatas)
 
         prompt = answer_with_context(context, query)
-        answer = gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0))
-        print(f'Asnwer : {answer}')
-        answer = f"{answer.answer}\n\n\nSource(s) :\n\n{result_metadatas}"
+        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 not st.session_state.documents_only:
             answer = f'Documents Based :\n\n{answer}'
@@ -255,15 +257,15 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
         print(f'Choice 2: {action}')
         if action == 'General Question':
             prompt = open_query_prompt(past_messages, query)
-            answer = gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)).answer
+            answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer']
         else:
             print(f'GLOBAL STATE : {global_state_documents_only}')
             if global_state_documents_only:
                 prompt = idk(query)
-                answer = gen_text(prompt, max_tokens=128, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)).answer
+                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 = gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)).answer
+                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()