diff --git a/scripts/custom_retriever.py b/scripts/custom_retriever.py index 27949cf..fce0a99 100644 --- a/scripts/custom_retriever.py +++ b/scripts/custom_retriever.py @@ -6,6 +6,7 @@ from llama_index.core import QueryBundle from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever from llama_index.core.schema import NodeWithScore, TextNode +from llama_index.postprocessor.cohere_rerank import CohereRerank logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) @@ -71,4 +72,8 @@ def filter_nodes_by_unique_doc_id(nodes): else: nodes_context.append(node) + reranker = CohereRerank(top_n=5, model="rerank-english-v3.0") + nodes_context = reranker.postprocess_nodes(nodes_context, query_bundle) + logfire.info(f"Cohere raranking to {len(nodes_context)} nodes") + return nodes_context diff --git a/scripts/gradio-ui.py b/scripts/gradio-ui.py index 8428507..152133e 100644 --- a/scripts/gradio-ui.py +++ b/scripts/gradio-ui.py @@ -1,3 +1,4 @@ +import logging import os import pickle @@ -22,6 +23,9 @@ load_dotenv() +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) +logging.getLogger("httpx").setLevel(logging.WARNING) logfire.configure() @@ -67,11 +71,10 @@ ] -# # Initialize MongoDB # mongo_db = ( # init_mongo_db(uri=MONGODB_URI, db_name="towardsai-buster") # if MONGODB_URI -# else logger.warning("No mongodb uri found, you will not be able to save data.") +# else logfire.warn("No mongodb uri found, you will not be able to save data.") # ) @@ -223,11 +226,40 @@ def generate_completion( def vote(data: gr.LikeData): + collection = "liked_data-test" if data.liked: print("You upvoted this response: " + data.value["value"]) else: print("You downvoted this response: " + data.value["value"]) + # completion_json["liked"] = like_data.liked + # logger.info(f"User reported {like_data.liked=}") + + # try: + # cfg.mongo_db[collection].insert_one(completion_json) + # except: + # logger.info("Something went wrong logging") + + +# def save_completion(completion: Completion, history): +# collection = "completion_data-hf" + +# # Convert completion to JSON and ignore certain columns +# completion_json = completion.to_json( +# columns_to_ignore=["embedding", "similarity", "similarity_to_answer"] +# ) + +# # Add the current date and time to the JSON +# completion_json["timestamp"] = datetime.utcnow().isoformat() +# completion_json["history"] = history +# completion_json["history_len"] = len(history) + +# try: +# cfg.mongo_db[collection].insert_one(completion_json) +# logger.info("Completion saved to db") +# except Exception as e: +# logger.info(f"Something went wrong logging completion to db: {e}") + accordion = gr.Accordion(label="Customize Sources (Click to expand)", open=False) sources = gr.CheckboxGroup(