지식 그래프 기반 추천 모델을 위한 데이터셋 구성 방법

📄 설명

Knowledge-aware 형식의 모델을 돌리기위해서는 .inter, .kg, .link 형식의 특별한 데이터 구조가 필요합니다. 변환하는 과정에서 겪은 문제와 해결 방법을 공유합니다.

그니까 요약하자면 데이터의 .kg, .link의 데이터가 왜 그렇게 구성되어 있냐면 옛날 데이터라 Wikidata가 아니라 옛날 구글이 쓰던 Freebase형식으로 구성되있던 시절 데이터라 그럼. 따로 옛날 Freebase형식에 맞게 수정이 필요.

✅ 작업할 내용

  • 지식 그래프의 특성 이해
  • test.link, test.inter, test.kg 등 베이스 데이터셋 분석
  • Freebase와 Wikidata에 대한 이해, 왜 데이터셋이 Freebase방식으로 작성되어있는지 의문점 해결
  • Freebase 형식으로 변환
import pandas as pd
import json
import os
from collections import defaultdict

def convert_to_freebase_format(input_dir, output_dir):
    """MovieLens 데이터를 Freebase 형식으로 변환"""
    os.makedirs(output_dir, exist_ok=True)
    print(f"입력 디렉토리: {input_dir}")
    print(f"출력 디렉토리: {output_dir}")
    
    # Freebase 스타일 ID 생성 함수
    def create_fb_id(prefix, id_num):
        return f"m.{prefix}{id_num:06x}"  # 16진수 6자리로 변환
    
    # 1. ID 매핑 생성
    ratings_df = pd.read_csv(os.path.join(input_dir, 'train_ratings.csv'))
    print(f"ratings 데이터 로드 완료: {len(ratings_df)} 행")
    
    # 기존 ID를 Freebase 스타일 ID로 매핑
    # int()를 사용하여 numpy.int64를 Python int로 변환
    item_to_fb = {int(item): create_fb_id('i', idx) 
                 for idx, item in enumerate(ratings_df['item'].unique())}
    user_to_fb = {int(user): int(idx) 
                 for idx, user in enumerate(ratings_df['user'].unique(), 1)}
    
    # 2. inter 파일 생성
    ratings_processed = pd.DataFrame({
        'user_id:token': ratings_df['user'].map(lambda x: user_to_fb[int(x)]),
        'item_id:token': ratings_df['item'].map(lambda x: int(x)),
        'rating:float': [1.0] * len(ratings_df),
        'timestamp:float': ratings_df['time']
    })
    
    # inter 파일 저장
    ratings_processed.to_csv(
        os.path.join(output_dir, 'movie.inter'),
        sep='\t',
        index=False
    )
    print("movie.inter 파일 저장 완료")
    
    # 3. Knowledge Graph (kg) 생성
    kg_triples = []
    
    # 감독 정보 처리
    directors_df = pd.read_csv(os.path.join(input_dir, 'directors.tsv'), sep='\t')
    for _, row in directors_df.iterrows():
        if int(row['item']) in item_to_fb:
            head = item_to_fb[int(row['item'])]
            relation = 'film.film.directed_by'
            tail = f"d.{row['director']}"
            kg_triples.append([head, relation, tail])
    print(f"감독 정보 처리 완료: {len(directors_df)} 행")
    
    # 장르 정보 처리
    genres_df = pd.read_csv(os.path.join(input_dir, 'genres.tsv'), sep='\t')
    for _, row in genres_df.iterrows():
        if int(row['item']) in item_to_fb:
            head = item_to_fb[int(row['item'])]
            relation = 'film.film.genre'
            tail = f"g.{row['genre'].lower().replace(' ', '_')}"
            kg_triples.append([head, relation, tail])
    print(f"장르 정보 처리 완료: {len(genres_df)} 행")
    
    # 작가 정보 처리
    writers_df = pd.read_csv(os.path.join(input_dir, 'writers.tsv'), sep='\t')
    for _, row in writers_df.iterrows():
        if int(row['item']) in item_to_fb:
            head = item_to_fb[int(row['item'])]
            relation = 'film.film.writer'
            tail = f"w.{row['writer']}"
            kg_triples.append([head, relation, tail])
    print(f"작가 정보 처리 완료: {len(writers_df)} 행")
    
    # 연도 정보 처리
    years_df = pd.read_csv(os.path.join(input_dir, 'years.tsv'), sep='\t')
    for _, row in years_df.iterrows():
        if int(row['item']) in item_to_fb:
            head = item_to_fb[int(row['item'])]
            relation = 'film.film.year'
            tail = f"y.{row['year']}"
            kg_triples.append([head, relation, tail])
    print(f"연도 정보 처리 완료: {len(years_df)} 행")
    
    # kg 파일 저장
    kg_df = pd.DataFrame(kg_triples, columns=['head_id:token', 'relation_id:token', 'tail_id:token'])
    kg_df.to_csv(os.path.join(output_dir, 'movie.kg'), sep='\t', index=False)
    print("movie.kg 파일 저장 완료")
    
    # 4. link 파일 생성
    link_data = [[int(item), fb_id] for item, fb_id in item_to_fb.items()]
    link_df = pd.DataFrame(link_data, columns=['item_id:token', 'entity_id:token'])
    link_df.to_csv(os.path.join(output_dir, 'movie.link'), sep='\t', index=False)
    print("movie.link 파일 저장 완료")
    
    # 매핑 정보 저장 - int64를 int로 변환
    mappings = {
        'user_mapping': {int(k): int(v) for k, v in user_to_fb.items()},
        'item_mapping': {int(k): str(v) for k, v in item_to_fb.items()}  # Freebase ID는 문자열로 저장
    }
    
    with open(os.path.join(output_dir, 'mapping_info.json'), 'w') as f:
        json.dump(mappings, f, indent=2)
    
    return mappings

# 사용 예시
if __name__ == "__main__":
    input_directory = "data/train"
    output_directory = "recbole_data"
    
    try:
        mappings = convert_to_freebase_format(input_directory, output_directory)
        print("\n=== 변환 완료 ===")
        print(f"총 사용자 수: {len(mappings['user_mapping'])}")
        print(f"총 아이템 수: {len(mappings['item_mapping'])}")
        
        # 결과 샘플 출력
        print("\n=== 변환된 데이터 샘플 ===")
        for file_name in ['movie.inter', 'movie.kg', 'movie.link']:
            print(f"\n{file_name} 샘플:")
            df = pd.read_csv(os.path.join(output_directory, file_name), sep='\t', nrows=5)
            print(df)
            
    except Exception as e:
        print(f"\n오류 발생: {str(e)}")
        import traceback
        print(traceback.format_exc())