K-ragrec: Knowledge Graph Retrieval-Augmented Generation for LLM-based

본 코드는 저희 ACL 논문인 'Knowledge Graph Retrieval-Augmented Generation for LLM-based Recommendation'을 위한 pytorch 구현체입니다.

• Python==3.9
• numpy==1.23.4
• torch==2.4.1
• cuda==11.8.89
• transformers==4.45.2
• networkx==2.8.7
• peft==0.12.0

저희는 세 가지 데이터셋(MovieLens-1M, MovieLens-20M 및 Amazon Book)을 제공합니다.

공간 제약으로 인해 MovieLens-1M만 처리했으며, KG와 데이터셋은 다음 링크에 업로드했습니다: https://drive.google.com/file/d/1MlEPkRj47WrdXECUiz5D6Ie1oMv4hKC9/view?usp=sharing. 다른 데이터셋의 경우 논문에서와 같이 처리할 수 있는 원본 데이터를 제공합니다.

저희는 지식 그래프를 처리하고 해당 지식 벡터 데이터베이스를 제공합니다.

MovieLens-1M으로 K-RagRec을 실행하는 방법 (임계값 사용)

학습(Training) 시:

python train.py --model_name graph_llm --llm_model_name 7b --llm_frozen True --dataset ml1m --batch_size 5 --gnn_model_name gt --gnn_num_layers 4 --adaptive_ratio 0.5 --sub_graph_numbers 3 --reranking_numbers 5--adaptive_ratio 5

평가(evaluation) 시:

python evaluate.py --model_name graph_llm --llm_model_name 7b --llm_frozen True --dataset ml1m --batch_size 5 --gnn_model_name gt --gnn_num_layers 4 --adaptive_ratio 0.5 --sub_graph_numbers 3 --reranking_numbers 5--adaptive_ratio 5

또는 다음 명령어를 실행할 수 있습니다:

bash run.sh

• K-RagRec을 다른 데이터셋에서 실행하려면, 먼저 해당 데이터셋들을 처리하고 데이터셋 이름을 변경해야 합니다. 저희 환경에는 3가지 데이터셋이 제공됩니다.
• llm_model_name(예: 8b)을 다른 LLM 모델에 맞게 변경하세요. - 저희 실험은 두 개의 NVIDIA A6000-48G GPU에서 수행되었습니다. GPU가 하나만 있는 경우, graph_llm_for_one_GPU를 사용해 주십시오.

감사합니다! / 谢谢! / ありがとう! / merci! / 감사! / Danke! / спасибо! / gracias! ...

@article{wang2025knowledge,
title={Knowledge Graph Retrieval-Augmented Generation for LLM-based Recommendation},
author={Wang, Shijie and Fan, Wenqi and Feng, Yue and Ma, Xinyu and Wang, Shuaiqiang and Yin, Dawei},
...}