Update main.py

This commit is contained in:
Anooyman 2024-03-24 15:52:12 +08:00
parent 8c81c222a9
commit 773c70687e

View File

@ -1,6 +1,3 @@
import os
from config.config import data_dir
from data_processing import Data_process
from pipeline import EmoLLMRAG
from util.llm import get_glm
from loguru import logger
@ -14,23 +11,6 @@ from loguru import logger
'''
def main(query, system_prompt=''):
logger.info(data_dir)
if not os.path.exists(data_dir):
os.mkdir(data_dir)
dp = Data_process()
vector_db = dp.load_vector_db()
docs, retriever = dp.retrieve(query, vector_db, k=10)
logger.info(f'Query: {query}')
logger.info("Retrieve results===============================")
for i, doc in enumerate(docs):
logger.info(doc)
passages,scores = dp.rerank(query, docs)
logger.info("After reranking===============================")
for i in range(len(scores)):
logger.info(passages[i])
logger.info(f'score: {str(scores[i])}')
if __name__ == "__main__":
query = """
我现在处于高三阶段感到非常迷茫和害怕我觉得自己从出生以来就是多余的没有必要存在于这个世界