Update app.py

app.py

@@ -215,6 +215,7 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
 
     regex = build_regex_from_schema(schema, r"[\n ]?")
     gen_text = outlines.generate.regex(llm, regex)
+    gen_text.format_sequence = lambda x: schema_object.parse_raw(x)
 
     gen_choice = outlines.generate.choice(llm, choices=['Yes', 'No'])
     prompt = route_llm(context, query)
@@ -231,7 +232,7 @@ 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.6, top_p=0.9))['answer']
+        answer = gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9))
         answer = f"{answer}\n\n\nSource(s) :\n\n{result_metadatas}"
             
         if not st.session_state.documents_only:
@@ -243,14 +244,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))['answer']
+            answer = gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9))
         else:
             if st.session_state.documents_only:
                 prompt = idk(query)
-                answer = gen_text(prompt, max_tokens=128, sampling_params=SamplingParams(temperature=0.6, top_p=0.9))['answer']
+                answer = gen_text(prompt, max_tokens=128, sampling_params=SamplingParams(temperature=0.6, top_p=0.9))
+                print(f'TYPE: {type(answer)})
             else:
                 prompt = self_knowledge(query)
-                answer = gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9))['answer']
+                answer = gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9))
                 answer = f'Internal Knowledge :\n\n{answer}'
 
     torch.cuda.empty_cache()