Update app.py

app.py

@@ -157,20 +157,12 @@ def idk(query: str):
 
 @outlines.prompt
 def self_knowledge(query: str):
-    """
+    """Answer the following question by using your own knowledge about the topic.
 
     Question: {{ query }}
     """
 
 def main(query: str, client: QdrantClient, collection_name: str, llm, dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding, past_messages: str):
-    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))
-    )
-
     prompt = build_initial_prompt(query)
     gen_text = outlines.generate.text(llm)
     
@@ -194,6 +186,14 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
         action = gen_choice(prompt, max_tokens=128, sampling_params=SamplingParams(temperature=0))
 
         if action == 'Yes':
+            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))
+            )
+            
             prompt = answer_with_context(context, query)
             answer = gen_text(prompt, max_tokens=500, sampling_params=SamplingParams(temperature=0.3))
             answer = f"{answer}\n\n\nSource(s) :\n\n{result_metadatas}"
@@ -202,57 +202,14 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
                 answer = f'Documents Based :\n\n{answer}'
         else:
             if st.session_state.documents_only:
-                
+                prompt = idk(query)
+                answer = gen_text(prompt, max_tokens=500, sampling_params=SamplingParams(temperature=0.3))
+            else:
+                prompt = self_knowledge(query)
+                answer = gen_text(prompt, max_tokens=500, sampling_params=SamplingParams(temperature=0.3))
+                answer = f'Internal Knowledge :\n\n{answer}'
 
     return answer
-            
-            
-    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}'
 
 def collect_files(conn, cursor, directory, pattern):
     array = []
@@ -548,20 +505,11 @@ def chunk_documents(texts: List[str], metadatas: List[dict], dense_model: Optimu
 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': "{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
         }