Tangledup-ai/lang-agent
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rag test changes

Browse Source
This commit is contained in:
goulustis
2026-01-05 20:49:38 +08:00
parent 99d3683f13
commit 2814b3dabc
1 changed files with 4 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
scripts/make_rag_database.py
View File

@@ -15,6 +15,7 @@ def main(save_path = "assets/xiaozhan_emb"):
df_cat = pd.read_csv(cat_f) df_cat = pd.read_csv(cat_f)
df_desc = pd.read_csv(desc_f) df_desc = pd.read_csv(desc_f)
df_desc = df_desc[df_desc["is_available"] == 't'].reset_index(drop=True)
id_desc_dic = {} id_desc_dic = {}
for _, (id, name, desc) in df_cat[["id", "name", "description"]].iterrows(): for _, (id, name, desc) in df_cat[["id", "name", "description"]].iterrows():
@@ -33,12 +34,6 @@ def main(save_path = "assets/xiaozhan_emb"):
embeddings = QwenEmbeddings( embeddings = QwenEmbeddings(
api_key=os.environ.get("ALI_API_KEY") api_key=os.environ.get("ALI_API_KEY")
) )
# embeddings = OpenAIEmbeddings(
# model="text-embedding-v4",
# api_key=os.environ.get("ALI_API_KEY"),
# base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
# )
# embeddings = OpenAIEmbeddings()
if not osp.exists(save_path): if not osp.exists(save_path):
# --- STEP 2: Create vector store --- # --- STEP 2: Create vector store ---
@@ -66,12 +61,12 @@ def main(save_path = "assets/xiaozhan_emb"):
# --- STEP 5: Use the retriever/QA chain on the loaded store --- # --- STEP 5: Use the retriever/QA chain on the loaded store ---
retriever = loaded_vectorstore.as_retriever(search_kwargs={ retriever = loaded_vectorstore.as_retriever(search_kwargs={
"k":3 "k":5
}) })
u = loaded_vectorstore.similarity_search("灯与尘", k=2) u = loaded_vectorstore.similarity_search("灯与尘", k=2)
res = retriever.invoke("灯与尘") res = retriever.invoke("野心心")
for doc in res: for doc in res:
print(doc) print(doc)