Update app.py

app.py

@@ -8,59 +8,55 @@ pc = Pinecone(api_key=os.getenv('PINECONE_API_KEY'))
 
 
 from llama_index.vector_stores.pinecone import PineconeVectorStore
-
 from llama_index.core import VectorStoreIndex
 from llama_index.core.response.pprint_utils import pprint_source_node
-
-
 from llama_index.llms.octoai import OctoAI
 
 octoai = OctoAI(
     token=os.getenv('OCTOML_KEY'),
     model="meta-llama-3-70b-instruct",
-    # messages=[
-    #     {"role": "system", "content": "You are a helpful assistant."},
-    #     {"role": "user", "content": prompt_af}
-    # ],
     max_tokens=512,
-    # presence_penalty=0,
     temperature=0.1,
-    # top_p=0.9,
 )
 
 
 from llama_index.core.memory import ChatMemoryBuffer
 
-# memory = ChatMemoryBuffer.from_defaults(token_limit=5000)
+import gradio as gr
+from io import StringIO
 
-# chat_engine = vindex.as_chat_engine(
-#     chat_mode="context",
-#     llm=octoai,
-#     memory=memory,
-#     system_prompt=(
-#         "You are a chatbot, able to have normal interactions, as well as talk about news events."
-#     ),
-# )
+def get_credit_dist(history):
+    atoms_l = cu.sentence_splitter.split_text(history[-1][1])
+    atoms_l = list(filter(lambda x: len(x) > 50, atoms_l))
+    atom_topkmatches_l = cu.get_atom_topk_matches_l_concurrent(atoms_l, max_workers=8)
 
+    atomidx_w_single_url_aggmatch_l = cu.aggregate_atom_topkmatches_l(atom_topkmatches_l)
+    atom_support_l = cu.get_atmom_support_l_from_atomidx_w_single_url_aggmatch_l_concurrent(atoms_l, atomidx_w_single_url_aggmatch_l, max_workers=8)
 
+    credit_dist = cu.credit_atom_support_list(atom_support_l)
 
-import gradio as gr
+    _out = StringIO()
+    print(f"Credit distribution to sources:\n", file=_out)
+    cu.print_credit_dist(credit_dist, prefix='    ', url_to_id=None, file=_out)
+    print(file=_out)
 
-# def predict(message, history):
-#     response = chat_engine.stream_chat(message)
+    print(f"Per claim support:\n", file=_out)
+    for j, atom_support in enumerate(atom_support_l):
+        print(f"    Claim {j+1}: \"{atoms_l[j]}\"\n", file=_out)
+        cu.print_atom_support(atom_support, prefix='            ', file=_out)
+        print(file=_out)
 
-#     partial_message = ""
-#     for token in response.response_gen:
-#         partial_message += token
-#         yield partial_message
+    return _out.getvalue()
 
-# demo = gr.ChatInterface(predict)
 
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
+    chatbot = gr.Chatbot(height=800)
     msg = gr.Textbox()
     clear = gr.Button("Clear")
 
+    credit_box = gr.Textbox(label="Credit distribution", lines=20, autoscroll=False)
+    credit_btn = gr.Button("Credit response")
+
     def get_chat_engine():
         vector_store = PineconeVectorStore(pinecone_index=pc.Index('prorata-postman-ds-256'))
         vindex = VectorStoreIndex.from_vector_store(vector_store)
@@ -87,6 +83,8 @@ with gr.Blocks() as demo:
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(bot, [chatbot, chat_engine_var], chatbot)
     clear.click(lambda x: x.reset(), chat_engine_var, chatbot, queue=False)
+
+    credit_btn.click(get_credit_dist, chatbot, credit_box)
     
 if __name__ == "__main__":
     demo.queue()