728x90
728x90
https://arxiv.org/abs/1705.03551
TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension
We present TriviaQA, a challenging reading comprehension dataset containing over 650K question-answer-evidence triples. TriviaQA includes 95K question-answer pairs authored by trivia enthusiasts and independently gathered evidence documents, six per questi
arxiv.org
이 데이터 셋을 사용할 것이기에...
적당히 봤습니다.
다른 데이터 셋들에 비해 발전된 데이터다!
질문이 문서 그대로가 아니라 조금씩 차이가 난다!
| 🧩 문제 상황 (Motivation) | 기존 RC 데이터셋은 • 질문과 문서가 너무 밀접하게 연결되어 있고 • 어휘 및 구문 표현이 단순하여 → 모델 성능이 과대평가되는 문제 발생 |
| 🎯 연구 목적 (Goal) | • 실제 사용되는 자연 질문 기반의 RC 데이터셋 구축 • 질문과 문서를 독립적으로 수집하여 어휘/구문 다양성 확보 • 다중 문장 추론 및 간접 정답 포함 등 더 어려운 QA 문제 구성 |
| 🏗️ 방법론 (Methodology) | • 14개 Trivia 사이트에서 자연 질문 96K개 수집 • Bing API + TAGME로 Wikipedia / Web 문서 수집 • 정답 문자열 포함 여부로 distant supervision 적용 • Web은 문서당 평가 / Wikipedia는 질문당 평가 |
| 📊 실험 결과 (Results) | • BiDAF 성능: EM 40%, Human: 79% • 기존 SQuAD보다 성능 낮음 → 난이도 높음 입증 • 오류 원인: 문서 누락, paraphrasing, 다중 문장 추론 실패 등 |
| 🧠 기여 (Contributions) | • 세계 최대 규모의 RC 데이터셋 중 하나 (650K triplets) • 질문-문서 독립 생성 → 표현 다양성 확보 • 멀티 문장 reasoning, world knowledge, paraphrasing 요구 • SOTA 모델도 해결 어려움 → 도전적 벤치마크 제시 • verified subset 제공으로 정확한 평가 가능 |
| ⚠️ 한계 (Limitations) | • distant supervision → 일부 정답 누락 가능 • BiDAF는 긴 문서 처리에 부적합 (800단어 제한) • 정답 위치 span 제공되지 않음 • 정답이 아예 존재하지 않는 경우 처리 안됨 |
더보기
🔍 1. 논문 개요
| 문제 | 기존 RC 데이터셋은 질문이 너무 단순하거나 문서와 질문의 연관성이 과도하게 높아, 모델 성능이 과대평가됨 |
| 목표 | 보다 자연스럽고, 복잡하며, 어려운 질문-문서-정답 세트를 포함한 데이터셋 구성 및 베이스라인 성능 분석 |
📌 2. 문제 정의와 TriviaQA의 목적
- 문제 정의: 주어진 질문 q에 대해, 하나 이상의 문서 D를 근거로 정답 a를 도출하는 Reading Comprehension 문제
- TriviaQA의 특징:
- 질문과 문서의 독립성 확보: 질문은 Trivia 애호가들이 출제, 문서는 이후 수집됨 → 편향 적음
- 멀티 문장 추론 필요: 단일 문장이 아닌 여러 문장 간 추론
- 구문 및 어휘 다양성 존재 → 단순한 패턴 매칭으로는 어렵도록 설계
📊 3. TriviaQA 데이터셋 구성
| 구성 요소 | 수치 및 설명 |
| 총 QA 쌍 | 약 96,000개 |
| 문서 수 | 약 662,000개 |
| 평균 문서 길이 | 2,895 단어 |
| 질문당 문서 수 | 평균 6개 문서 |
| 데이터 출처 | Web 문서 + Wikipedia |
| 정답 유형 | Wikipedia 제목(93%), 수치(4%), 자유 텍스트(3%) 등 |
| 다중 문장 추론 필요 | 40% (SQuAD 대비 약 3배) |
🔍 4. 데이터 수집 과정
- 질문 수집: 14개의 Trivia 웹사이트에서 자연스럽게 생성된 질문 수집 (길이 ≥ 4단어)
- 문서 수집:
- Web: Bing 검색 API를 통해 상위 10개 URL 크롤링 (PDF, trivia 웹 제외)
- Wikipedia: TAGME로 엔티티 연결 후 해당 페이지 수집
- 정답 포함 문서 필터링: 정답이 포함된 문서만 유지하여 학습을 위한 distant supervision 제공
🧠 5. 난이도 분석 및 질적 분석
| 분류 | 설명 | 빈도 (Wikipedia 기준) |
| 구문적 변화 | 질문과 문장이 구문적으로 다름 | 69% |
| 어휘적 변화 | 동의어/유사 표현 존재 | 41% |
| 배경지식 필요 | 상식/지식 기반 추론 필요 | 17% |
| 다중 문장 추론 | 여러 문장 조합 필요 | 40% |
| 리스트/표 기반 | 리스트/표 안에 정답 존재 | 7% (Web 기준) |
🧪 6. 실험 및 베이스라인
🧩 실험 구성
- 두 영역으로 평가: Wikipedia / Web
- 3가지 모델:
- Random Entity Baseline (랜덤 추측)
- Entity Classifier (LambdaMART 기반)
- BiDAF (Neural RC 모델)
📈 성능 비교 (Wikipedia domain)
| 모델 | EM | F1 | Oracle | Human |
| Random | 12.7 | 22.9 | 16.3 | - |
| Classifier | 22.4 | 26.5 | 71.6 | - |
| BiDAF | 40.3 | 45.9 | 82.8 | 79.7 |
→ BiDAF도 인간보다 훨씬 낮은 성능 → 여전히 도전적인 문제임을 입증
❌ 7. 주요 에러 유형
| 에러 유형 | 비율 (%) |
| 문서 내 정답 없음 | 19 |
| 문서 잘림으로 인한 손실 | 15 |
| 패러프레이징 이해 부족 | 29 |
| 정답 유사 표현(낚시) | 11 |
| 다중 문장 추론 실패 | 18 |
→ 대부분 RC 모델의 한계점인 간접 표현, 문서 길이, 다중 문장 reasoning과 관련
✅ 8. 논문의 기여
| 질문과 문서를 독립적으로 수집하여 더 현실적인 QA 문제 구성 |
| 650K 이상의 대규모 QA-Evidence 트리플 구축 |
| 기계 독해가 어려운 현상 (어휘/구문 다양성, reasoning, 다중 문장)에 대한 체계적 분석 |
| baseline 결과 제시 → 현재 모델로 해결 어려움을 보여줌 |
| Web과 Wikipedia 도메인 모두 지원하는 멀티 도메인 QA 벤치마크 |
📌 한 줄 요약
TriviaQA는 자연적으로 생성된 질문과 문서를 바탕으로 복잡한 reasoning과 높은 어휘/구문 다양성을 요구하는 대규모 QA 벤치마크로, 기존 RC 모델에 도전적인 평가 환경을 제공한다.
📍 1. TriviaQA를 만든 이유 (Why)
✅ 기존 데이터셋의 한계
- 단순하고 인위적인 질문 구성: SQuAD, NewsQA 등의 기존 데이터셋은 질문이 문서에서 파생되거나, 문서와 질문이 너무 강하게 연관되어 있음.
- 낮은 어휘/구문 다양성: 질문과 문서 간의 표현 차이가 적어, 모델이 단순한 패턴 학습만으로도 높은 성능을 달성 가능함.
- 단일 문장 추론에 국한: 대부분의 질문은 하나의 문장에서 정답을 찾을 수 있음 → 다중 문장(reasoning across sentences) 처리가 필요 없는 경우가 많음.
- 빠른 포화 상태: 새로운 데이터셋이 나올 때마다 모델 성능은 급속히 향상되어 곧 인간 수준을 초과함 → 더 어려운 벤치마크가 필요
✅ 목표
보다 현실적이고 어려운 reading comprehension 문제를 만들자.
- 실제 세계에서처럼 질문과 문서가 독립적으로 생성된 환경
- 어휘 및 구문 표현이 다양하고 비정형적
- 정답을 찾기 위해 여러 문장을 통합해 추론해야 하는 복잡한 문제 제공
🎯 2. TriviaQA의 목적 (Goal)
| 자연스러운 질문 구성 | 실제 Trivia 팬들이 만든 퀴즈 기반 질문을 사용하여, 과도하게 단순화된 RC 문제를 지양 |
| 문서 다양성 | Web 검색 결과와 Wikipedia에서 문서를 수집하여, 문체, 형식, 정보의 다양성 제공 |
| 복잡한 reasoning 유도 | 정답을 찾기 위해 여러 문장을 통합하거나, 배경지식, 시간/비교 reasoning이 필요하도록 설계 |
| RC 모델의 한계 평가 | 기존 SOTA 모델(BiDAF 등)이 어디서 실패하는지 파악하여 향후 연구 방향 제시 |
| 다양한 QA 태스크 적용 가능 | Open-domain QA, IR QA, KB QA 등 다양한 방식의 QA 연구에 활용 가능하게 설계됨 |
🏗️ 3. 데이터셋 생성 과정 (Construction Process)
📌 3.1 질문 수집 (Question Collection)
- 출처: 14개의 Trivia 및 Quiz League 웹사이트
- 형태: 자연어로 완전하게 기술된 문장형 질문 (예: “Who directed the film Inception?”)
- 필터링: 4단어 이하의 질문은 제거 → 너무 단순하거나 의미 불분명한 질문 배제
- 총 질문 수: 약 96,000개
📌 3.2 문서 수집 (Evidence Collection)
✅ (1) Web Evidence
- 방법: 각 질문을 Bing Web 검색 API에 쿼리로 넣고, 상위 50개의 URL 수집
- 후처리:
- PDF 및 비정형 문서 제거
- ‘trivia’, ‘question’, ‘answer’ 등의 단어가 포함된 URL 제거 (출처 편향 방지)
- 최종적으로 상위 10개의 HTML 페이지를 수집 및 정제
- 문서 형태: 블로그, 뉴스, 위키, 잡지 등 다양한 문서 스타일
✅ (2) Wikipedia Evidence
- 방법:
- 질문 내의 명사구(엔티티)를 TAGME 엔티티 링크 툴로 추출
- 매칭된 Wikipedia 문서들을 추가적인 evidence로 수집
- 평균 문서 수: 질문당 1.8개의 위키 문서
📌 3.3 정답 포함 여부 필터링 (Distant Supervision)
- Distant Supervision 방식 사용
- 정답 문자열이 문서 안에 포함되어 있으면 해당 문서를 evidence로 간주
- 정답이 없는 문서는 제거
- Web domain: 정답 포함된 각 문서를 개별 QA 트리플로 구성 (많은 중복 있음)
- Wikipedia domain: 질문별로 관련 문서를 풀링(pooling) → 질문 하나에 여러 문서가 연결됨
📌 3.4 정제 및 분할
| 총 QA-Evidence 쌍 | 약 650,000개 |
| QA 쌍 수 | 약 95,956개 |
| 문서 수 | 약 662,659개 |
| 평균 질문 길이 | 14단어 |
| 데이터 분할 | train/dev/test = 80/10/10 |
| 추가 제공 | 1975개의 인간 검증 verified 세트 포함 (정답 존재 보장) |
✅ 정리
TriviaQA는 단순히 “질문과 정답”만 묶은 데이터셋이 아니라, 실제 인간이 만든 복잡한 질문, 다양한 스타일의 비구조 문서, 그리고 간접 학습을 위한 distant supervision을 활용해 생성된 고난이도 RC 벤치마크입니다. 다음과 같은 목표를 위해 설계되었습니다:
- 복잡한 reasoning 능력 평가
- 문서 다양성과 노이즈에 대한 모델의 견고성 평가
- 단순 정보 매칭을 넘어선 진짜 독해 능력 요구
📊 1. 결과 (Results)
✅ 평가 방식
- 도메인 분리:
- Wikipedia: 정답이 드물게 나타남 → 질문 단위 평가
- Web: 정답이 여러 문서에 반복됨 → 문서 단위 평가
- 평가지표: Exact Match (EM), F1 Score
- 비교 모델:
- Random Entity Baseline
- Entity Classifier (LambdaMART)
- BiDAF (Neural RC 모델)
📈 성능 요약
| 모델 | 도메인 | EM (Dev/Test) | F1 (Dev/Test) | Oracle | Human |
| Random | Wiki | 12.7 / 12.7 | 22.9 / 22.4 | ~16.3 | - |
| Classifier | Wiki | 23.4 / 22.4 | 27.6 / 26.5 | ~71.6 | - |
| BiDAF | Wiki | 40.3 / 40.3 | 45.9 / 45.9 | ~82.8 | 79.7 |
| Classifier | Web | 24.6 / 24.0 | 29.1 / 28.3 | ~66.3 | - |
| BiDAF | Web | 41.1 / 40.7 | 47.4 / 47.0 | ~83.0 | 75.4 |
📌 해석:
- BiDAF가 가장 높은 성능을 보였지만, 여전히 인간 성능(79~80%)에는 크게 못 미침
- 기존 SQuAD 성능(68%) 대비 TriviaQA 성능(40%)은 훨씬 낮음 → 훨씬 어려운 데이터셋임을 입증
❌ 오류 분석 (Error Analysis)
| 주요 오류 유형 | 설명 | 비율 (%) |
| 문서에 정답 없음 | 검색된 문서들에 정답 관련 내용 부재 | 19% |
| 문서 잘림으로 인한 정보 손실 | 길이 제한으로 정답 포함된 부분 누락 | 15% |
| 패러프레이징 실패 | 질문과 문서 간 표현 불일치 | 29% |
| 유사 엔티티 낚시(distractors) | 비슷한 단어가 많아 오답 선택 | 11% |
| 다중 문장 추론 실패 | 여러 문장을 결합해 추론 필요 | 18% |
🧾 2. 결론 (Conclusion)
- TriviaQA는 복잡하고 현실적인 Reading Comprehension 문제를 제공하는 데이터셋이다.
- 질문은 사람이 실제로 작성한 문장형 질문이며, 문서는 Web/Wikipedia에서 자동 수집
- 기존 SOTA 모델(BiDAF)도 TriviaQA에서는 인간 수준에 크게 미치지 못함 (40% vs 80%)
- 따라서 TriviaQA는 기계 독해(MRC)의 진정한 한계를 평가하고, 발전 방향을 제시하는 벤치마크로 적합
⚠️ 3. 한계 (Limitations)
| 한계 | 설명 |
| Distant supervision 의존 | 정답 문자열이 문서에 포함된 것만 선택 → 일부 문서에 정답이 암시적인 경우 누락 |
| 정답 위치 정확도 미흡 | BiDAF는 SQuAD 기반이라 span 추출 필요하지만, TriviaQA는 정답 위치를 명시하지 않음 |
| 긴 문서 처리 어려움 | BiDAF는 평균 122단어 기준 학습되어 있어, TriviaQA의 평균 2,800단어 문서에 부적합함 |
| 일부 noise 존재 | Web 문서에는 낚시 엔티티, 불완전한 정보 포함 가능성 있음 |
| 추론 기반 정답이 없는 경우 미처리 | 일부 질문은 정답이 직접적으로 나오지 않음 → “없음”을 예측하는 설정은 고려되지 않음 |
🏅 4. 주요 기여 (Contributions)
| ✅ 대규모 자연 질문 기반 데이터셋 구축 | 65만 개 QA-evidence 트리플, 질문은 trivia 팬들이 실제로 작성 |
| ✅ 문서-질문 독립 수집 | 기존 SQuAD와 다르게, 질문과 문서를 분리해 수집 → 표현 다양성 및 현실성 확보 |
| ✅ 어휘/구문 다양성, reasoning 난이도 높음 | 다중 문장 reasoning 필요, paraphrasing, world knowledge 포함 |
| ✅ SOTA 모델로 평가한 어려움 입증 | BiDAF도 인간 수준보다 훨씬 낮은 정확도 → 실제 난이도 있는 벤치마크 |
| ✅ 다양한 QA 연구에 활용 가능 | IR QA, Open-domain QA, KB QA 등 다양한 QA 패러다임 적용 가능 |
| ✅ verified subset 제공 | 인간 검증된 정답 포함 세트(1975개) 제공 → 정확한 모델 평가 가능 |
✨ 정리 한 줄 요약
TriviaQA는 질문과 문서가 독립적으로 생성되고, 어휘/구문 다양성과 멀티 문장 추론을 요구하며, 현재 SOTA 모델이 해결하지 못하는 고난이도 기계 독해 벤치마크로서 큰 가치를 지닌다.
https://github.com/mandarjoshi90/triviaqa
GitHub - mandarjoshi90/triviaqa: Code for the TriviaQA reading comprehension dataset
Code for the TriviaQA reading comprehension dataset - mandarjoshi90/triviaqa
github.com
https://huggingface.co/datasets/mandarjoshi/trivia_qa
mandarjoshi/trivia_qa · Datasets at Hugging Face
⚖️ CATIE-AQ/Guide_Evaluation_LLM 🚀 PankhuriSharma9795/SHL_Recommend_System 🚀 PankhuriSharma9795/SHL-Recommend
huggingface.co
일단 원본 데이터 셋은 여기 있습니다.
Gen Read
더보기
GenRead: 검색보다는 생성하자!
개요
GenRead는 "Generate rather than Retrieve: Large Language Models are Strong Context Generators" 논문의 공식 구현체입니다. 이 연구는 ICLR 2023에 발표되었으며, 기존의 RAG(Retrieve-Augmented Generation) 방식과 달리 외부 지식 베이스에서 문서를 검색하는 대신 Large Language Model(LLM)이 직접 배경 문서를 생성하여 질문에 답변하는 새로운 접근법을 제안합니다.
핵심 아이디어
- 기존 RAG: 질문 → 문서 검색 → 검색된 문서 기반 답변 생성
- GenRead: 질문 → 배경 문서 생성 → 생성된 문서 기반 답변 생성
시스템 아키텍처
GenRead는 2단계 파이프라인으로 구성됩니다:
Step 1: 배경 문서 생성 (Background Document Generation)
- 입력: 질문(query)
- 처리: LLM이 질문과 관련된 Wikipedia 스타일의 배경 문서를 생성
- 출력: 생성된 배경 문서
Step 2: 답변 추출 (Answer Extraction)
- 입력: 생성된 배경 문서 + 원본 질문
- 처리: 배경 문서를 참조하여 질문에 대한 답변 추출
- 출력: 최종 답변
프로젝트 구조
GenRead/
├── mainfunc.py # 메인 실행 파일
├── inference.py # OpenAI API 호출 및 추론 로직
├── evaluation.py # 평가 메트릭 계산
├── clusterfunc.py # 클러스터링 기반 문서 생성
├── indatasets/ # 입력 데이터셋
│ ├── webq/
│ └── fm2/
├── inprompts/ # 프롬프트 템플릿
│ ├── regular.jsonl
│ └── cluster.jsonl
└── README_KOR.md # 한국어 문서
데이터셋 형식
입력 데이터 형식 (indatasets/{dataset}/{dataset}-{split}.jsonl)
{
"id": "3267",
"question": "where was ancient carthage located?",
"answer": ["Tunisia"]
}
중간 결과 형식 (Step 1 출력)
{
"question": "where was ancient carthage located?",
"answer": ["Tunisia"],
"output": ["Ancient Carthage was located in present-day Tunisia..."]
}
프롬프트 시스템
Step 1 프롬프트 (배경 문서 생성)
Provide a background document from Wikipedia to answer the given question.
{query}
Step 2 프롬프트 (답변 추출)
Refer to the passage below and answer the following question with just one entity.
Passage: {background}
Question: {query}
The answer is
지원하는 모델
1. OpenAI 모델
- 지원 모델: GPT-4, GPT-4-turbo, GPT-3.5-turbo 등 (모델명에 'gpt' 포함)
- 처리 방식: 최신 OpenAI API v1.0+ 사용, 멀티스레드 기반 병렬 API 호출
- API 타입: GPT 모델은 Chat Completions API, 레거시 모델은 Completions API 자동 선택
- 장점: 높은 품질, 다양한 모델 선택, 최신 API 기능 지원
- 단점: API 비용 발생
2. HuggingFace 모델 (VLLM)
- 지원 모델: Llama, Mistral, Qwen 등 오픈소스 모델
- 처리 방식: VLLM을 통한 배치 추론
- GPU 관리: 텐서 병렬화 및 메모리 사용률 제어
- 장점: 비용 효율적, 높은 처리 속도, 다중 GPU 지원
- 단점: GPU 메모리 요구사항
GPU 자원 관리
- 텐서 병렬화: 큰 모델을 여러 GPU에 분산
- 메모리 관리: GPU 메모리 사용률 조절로 다른 프로세스와 공존
- GPU 선택: 특정 GPU만 사용하여 서버 자원 충돌 방지
지원하는 태스크 유형
1. Question Answering (질문 답변)
- 데이터셋: NQ, TriviaQA, WebQ, TWiki
- 평가 메트릭:
- Step 1: Recall@K (생성된 문서에 정답이 포함되어 있는가)
- Step 2: Exact Match (정확한 답변 일치도)
2. Fact Checking (사실 검증)
- 데이터셋: FEVER, FM2
- 평가 메트릭: Accuracy (참/거짓 분류 정확도)
3. Dialogue System (대화 시스템)
- 데이터셋: Wizard
- 평가 메트릭: F1-score, Rouge-L
설치 및 설정
1. 환경 설정
pip install -r requirements.txt
선택적 의존성
- VLLM (HuggingFace 모델용): pip install vllm
- OpenAI API: OpenAI 계정 및 API 키 필요
2. OpenAI API 키 설정
환경변수 설정 (권장)
export OPENAI_API_KEY="your_openai_api_key_here"
또는 코드에서 직접 설정
inference.py 파일의 26번째 줄에서 API 키를 설정할 수 있습니다:
client = OpenAI(
api_key="your_openai_api_key_here"
)
3. 데이터셋 다운로드
데이터셋을 indatasets/ 폴더에 배치하세요.
사용 방법
Zero-shot 설정
Step 1: 배경 문서 생성
OpenAI 모델 사용
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine gpt-4 \
--max_workers 10
HuggingFace 모델 사용 (VLLM)
# 단일 GPU 사용
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine meta-llama/Llama-2-7b-chat-hf \
--tensor_parallel_size 1 \
--gpu_memory_utilization 0.8 \
--cuda_visible_devices 0
# 다중 GPU 사용 (텐서 병렬화)
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine meta-llama/Llama-2-13b-chat-hf \
--tensor_parallel_size 2 \
--gpu_memory_utilization 0.9 \
--cuda_visible_devices 0,1
파라미터 설명:
- --dataset: 사용할 데이터셋 (webq, nq, tqa, fever, fm2, wizard)
- --task: 실행할 단계 (step1 또는 step2)
- --split: 데이터 분할 (train, dev, test)
- --engine: 사용할 모델 (OpenAI: gpt-*, HuggingFace: 모델 경로)
- --max_workers: OpenAI API 동시 요청 수 (기본값: 5)
- --tensor_parallel_size: VLLM 텐서 병렬화용 GPU 수 (기본값: 1)
- --gpu_memory_utilization: VLLM GPU 메모리 사용률 (0.0-1.0, 기본값: 0.9)
- --cuda_visible_devices: 사용할 GPU ID들 (예: "0,1,2")
Step 2: 답변 추출
python mainfunc.py \
--dataset webq \
--task step2 \
--split test \
--engine gpt-4 \
--max_workers 10
사용 예시
1. OpenAI GPT-4를 사용한 고품질 생성
# Step 1: 배경 문서 생성 (10개 워커로 빠른 처리)
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine gpt-4 \
--max_workers 10
# Step 2: 답변 추출
python mainfunc.py \
--dataset webq \
--task step2 \
--split test \
--engine gpt-4 \
--max_workers 10
2. Llama-2를 사용한 비용 효율적 생성 (단일 GPU)
# Step 1: 배경 문서 생성 (GPU 0 사용, 메모리 80% 활용)
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine meta-llama/Llama-2-13b-chat-hf \
--tensor_parallel_size 1 \
--gpu_memory_utilization 0.8 \
--cuda_visible_devices 0 \
--temperature 0.7
# Step 2: 답변 추출
python mainfunc.py \
--dataset webq \
--task step2 \
--split test \
--engine meta-llama/Llama-2-13b-chat-hf \
--tensor_parallel_size 1 \
--gpu_memory_utilization 0.8 \
--cuda_visible_devices 0
3. 대형 모델 다중 GPU 활용
# Llama-2 70B 모델을 4개 GPU에 분산
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine meta-llama/Llama-2-70b-chat-hf \
--tensor_parallel_size 4 \
--gpu_memory_utilization 0.9 \
--cuda_visible_devices 0,1,2,3
3. 다중 시퀀스 생성 (Supervised 설정)
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine gpt-4-turbo \
--num_sequence 5 \
--temperature 0.8 \
--max_workers 8
4. 공유 서버 환경에서의 GPU 자원 관리
# GPU 2,3번만 사용하여 다른 사용자와 충돌 방지
python mainfunc.py \
--dataset webq \
--task step1 \
--split test \
--engine meta-llama/Llama-2-13b-chat-hf \
--tensor_parallel_size 2 \
--gpu_memory_utilization 0.7 \
--cuda_visible_devices 2,3
코드 동작 원리
1. 모델 타입 자동 감지 (inference.py)
is_openai_model() 함수
def is_openai_model(engine):
"""Check if the model is an OpenAI model by looking for 'gpt' in the name."""
return 'gpt' in engine.lower()
2. 추론 파이프라인
OpenAI API 멀티스레드 처리
def run_openai_inference_batch(inputs_with_prompts, engine, max_tokens,
num_sequence=1, temp=0, max_workers=5):
"""Run OpenAI inference with multithreading for batch processing."""
all_outputs = []
with ThreadPoolExecutor(max_workers=max_workers) as executor:
# Submit all tasks
future_to_index = {}
for i, prompt in enumerate(inputs_with_prompts):
future = executor.submit(run_openai_inference_single,
prompt, engine, max_tokens, num_sequence, temp)
future_to_index[future] = i
# Collect results in order
results = [None] * len(inputs_with_prompts)
for future in as_completed(future_to_index):
index = future_to_index[future]
results[index] = future.result()
return results
VLLM 배치 처리
def run_vllm_inference(inputs_with_prompts, engine, max_tokens,
num_sequence=1, temp=0):
"""Run inference using VLLM for HuggingFace models."""
# Initialize VLLM model
llm = LLM(model=engine, tensor_parallel_size=1)
# Set sampling parameters
sampling_params = SamplingParams(
temperature=temp,
max_tokens=max_tokens,
n=num_sequence
)
# Generate responses in batch
outputs = llm.generate(inputs_with_prompts, sampling_params)
return [completion.text for output in outputs for completion in output.outputs]
3. 메인 추론 함수
def run_inference(inputs_with_prompts, engine, max_tokens,
num_sequence=1, temp=0, max_workers=5):
"""Main inference function that routes to appropriate backend."""
if is_openai_model(engine):
print(f"Using OpenAI API with {max_workers} workers for model: {engine}")
return run_openai_inference_batch(inputs_with_prompts, engine,
max_tokens, num_sequence, temp, max_workers)
else:
print(f"Using VLLM for HuggingFace model: {engine}")
return run_vllm_inference(inputs_with_prompts, engine,
max_tokens, num_sequence, temp)
4. 배치 처리 및 데이터 관리
run_main() 함수
- 배치 단위로 데이터를 처리 (기본 배치 크기: 20)
- 재시작 지원: 기존 출력 파일이 있으면 이어서 처리
- 진행 상황 모니터링 (tqdm 프로그레스 바)
add_prompt() 함수
def add_prompt(item, prompt):
query = item['question']
prompt = prompt.replace('{query}', query)
if item.get('output'): # Step 2의 경우
backinfo = rmreturn(item['output'][0])
prompt = prompt.replace('{background}', backinfo)
return prompt
5. 평가 시스템 (evaluation.py)
Recall@K 평가 (Step 1)
tokenizer = SimpleTokenizer()
lines = open(infile, 'r').readlines()[1:]
has_answer_count = 0
for line in lines:
line = json.loads(line)
answer = line['answer']
output = ' || '.join(line['output'])
if has_answer(answer, output, tokenizer):
has_answer_count += 1
recall = has_answer_count / len(lines)
return recall
#### Exact Match 평가 (Step 2)
```python
def eval_question_answering(infile):
lines = open(infile, 'r').readlines()[1:]
exact_match_count = 0
for line in lines:
line = json.loads(line)
answer = line['answer']
output = line['output'][0]
if ems(output, answer):
exact_match_count += 1
em = exact_match_count / len(lines)
return em
6. 메인 컨트롤러 (mainfunc.py)
태스크 타입 자동 감지
if args.dataset in ['nq', 'webq', 'tqa', 'twiki']:
datatype = 'question answering'
elif args.dataset in ['fever', 'fm2']:
datatype = 'fact checking'
elif args.dataset in ['wizard']:
datatype = 'dialogue system'
토큰 길이 자동 설정
if args.task == 'step1':
max_tokens = 300 # 배경 문서 생성용
elif args.task == 'step2':
if datatype == 'dialogue system':
max_tokens = 50 # 대화 응답용
else:
max_tokens = 10 # 단답형 답변용
출력 형식
폴더 구조 (모델명에 따라 자동 생성)
# OpenAI 모델의 경우
backgrounds-greedy-gpt-3.5-turbo/webq/webq-test-p1.jsonl
finaloutput-greedy-gpt-3.5-turbo/webq/webq-test-p1.jsonl
# HuggingFace 모델의 경우 (슬래시를 언더스코어로 변환)
backgrounds-greedy-meta-llama_Llama-2-7b-chat-hf/webq/webq-test-p1.jsonl
finaloutput-greedy-meta-llama_Llama-2-7b-chat-hf/webq/webq-test-p1.jsonl
평가 결과 출력
# Recall@K 결과 (Step 1)
backgrounds-greedy-{모델명}/webq/webq-recall@k.jsonl
# 최종 평가 결과 (Step 2)
finaloutput-greedy-{모델명}/webq/webq-metrics.jsonl
클러스터링 기능
clusterfunc.py는 더 다양한 배경 문서를 생성하기 위한 고급 기능을 제공합니다:
- 임베딩 기반 클러스터링: 질문들을 의미적으로 유사한 그룹으로 클러스터링
- 클러스터별 프롬프트: 각 클러스터에 특화된 few-shot 예시 제공
- 다양성 증대: 더 다양한 관점의 배경 문서 생성
성능 최적화
OpenAI API 최적화
- 멀티스레드 처리: --max_workers 파라미터로 동시 요청 수 조절
- 자동 재시도: API 실패 시 최대 200회 재시도
- 타임아웃 처리: 20초 타임아웃으로 응답 없는 요청 방지
- 비용 관리: 배치 크기와 워커 수를 조절하여 API 사용량 최적화
VLLM 최적화
- 배치 처리: 모든 입력을 한 번에 처리하여 GPU 활용률 극대화
- 텐서 병렬화: --tensor_parallel_size로 모델을 여러 GPU에 분산
- 메모리 관리: --gpu_memory_utilization로 GPU 메모리 사용량 조절
- GPU 선택: --cuda_visible_devices로 특정 GPU만 사용
- 모델 캐싱: 동일 모델 재사용 시 로딩 시간 단축
GPU 설정 가이드라인
- 7B 모델: 1개 GPU, 메모리 사용률 0.8-0.9
- 13B 모델: 1-2개 GPU, 메모리 사용률 0.8-0.9
- 30B 모델: 2-4개 GPU, 메모리 사용률 0.9
- 70B 모델: 4-8개 GPU, 메모리 사용률 0.9
일반적인 최적화
- 진행 상황 저장: 중단된 작업을 이어서 실행 가능
- 배치 크기 조절: 메모리와 처리 속도의 균형점 찾기
확장성
새로운 데이터셋 추가
- indatasets/{dataset}/ 폴더 생성
- JSONL 형식으로 데이터 준비
- mainfunc.py에서 데이터 타입 매핑 추가
새로운 평가 메트릭 추가
- evaluation.py에 새로운 평가 함수 구현
- mainfunc.py에서 해당 함수 호출 로직 추가
주의사항 및 제한사항
OpenAI API 관련
- API 비용: 사용량에 따른 비용 발생 (특히 GPT-4는 고비용)
- Rate Limit: OpenAI API의 분당 요청 제한 고려 필요
- 재현성: 샘플링 기반 생성은 결과가 매번 다를 수 있음
VLLM 관련
- GPU 메모리: 대형 모델은 상당한 GPU 메모리 필요
- 모델 호환성: 일부 모델은 VLLM에서 지원되지 않을 수 있음
- 초기 로딩: 모델 첫 로딩 시 시간이 오래 걸릴 수 있음
- GPU 충돌: 공유 서버에서는 적절한 GPU 자원 할당 필요
- 메모리 단편화: 장시간 사용 시 GPU 메모리 재시작 필요할 수 있음
일반적인 주의사항
- 데이터 형식: 입력 데이터가 지정된 JSONL 형식을 준수해야 함
- 디스크 공간: 대량 데이터 처리 시 충분한 저장 공간 필요
- 네트워크: OpenAI API 사용 시 안정적인 인터넷 연결 필요
GRG
더보기
GRG (Generator-Retriever-Generator) 프로젝트
목차
- 소개
- 요구사항
- 설치 및 설정
- 프로젝트 구조
- Document Generator (DG)
- Document Generator Retriever (DGR)
- 사용법
- 데이터셋
- 인용
소개
Generator-Retriever-Generator: A Novel Approach to Open-domain Question Answering 논문을 위한 저장소입니다. 이 연구에서는 개방형 도메인 질문 응답 문제를 해결하기 위한 GRG 접근법을 제시합니다.
GRG 접근법
GRG는 다음과 같은 3단계 파이프라인으로 구성됩니다:
- Generator (DG): Few-shot 학습을 통해 질문과 관련된 문서를 생성
- Retriever (DGR): 생성된 문서에서 관련 정보를 검색
- Generator: 검색된 정보를 바탕으로 최종 답변 생성
요구사항
- Python 3.8+
- PyTorch 2.0+
- Transformers 4.30+
- LlamaIndex 0.6+
- OpenAI API 키 (Document Generator용)
설치 및 설정
- 환경 설정
$ conda create -n grg python=3.8
$ conda activate grg
$ pip install -r requirements.txt
- OpenAI API 키 설정
- DG/utils.py 파일에서 openai.api_key에 본인의 OpenAI API 키를 추가하세요.
프로젝트 구조
GRG/
├── DG/ # Document Generator 모듈
│ ├── indatasets/ # 입력 데이터셋 (NQ, TQA, WEBQ)
│ ├── inprompts/ # 입력 프롬프트 (NQ, TQA, WEBQ)
│ ├── outdataset/ # 출력 데이터셋
│ ├── main.py # 메인 실행 파일
│ ├── process.py # 처리 유틸리티
│ ├── utils.py # 유틸리티 함수
│ └── README.md
├── DGR/ # Document Generator Retriever 모듈
│ ├── indatasets/ # 입력 데이터셋
│ ├── outdatasets/ # 출력 데이터셋
│ ├── main.py # 메인 실행 파일
│ └── README.md
├── requirements.txt # 패키지 의존성
└── README.md # 프로젝트 설명서
Document Generator (DG)
Document Generator는 Few-shot 학습 방법론을 사용하여 질문과 관련된 문서를 생성하는 모듈입니다.
주요 기능
- OpenAI GPT 모델 (text-davinci-003, text-davinci-002)을 사용한 문서 생성
- 다양한 데이터셋 지원 (Natural Questions, TriviaQA, WebQuestions)
- 커스터마이징 가능한 생성 파라미터
사용법
cd DG
python main.py \
--dataset [nq, tqa, webq] \
--split [train, dev, test] \
--engine [text-davinci-003, text-davinci-002] \
--num_sequence [10, 50] \
--temperature 0.5 \
--max_tokens [100, 200, 300]
파라미터 설명
- --dataset: 사용할 데이터셋 (nq: Natural Questions, tqa: TriviaQA, webq: WebQuestions)
- --split: 데이터 분할 (train, dev, test)
- --engine: 사용할 OpenAI 모델
- --num_sequence: 생성할 시퀀스 수
- --temperature: 생성 온도 (0.0-1.0)
- --max_tokens: 최대 토큰 수
Document Generator Retriever (DGR)
Document Generator Retriever는 생성된 문서에서 관련 정보를 검색하기 위해 특별히 설계된 검색기로, Sentence Transformers를 활용하여 성능을 향상시킵니다.
주요 기능
- Sentence Transformers 기반 임베딩
- 벡터 기반 유사도 검색
- 커스터마이징 가능한 top-k 검색
- 다양한 사전 훈련된 모델 지원
사용법
cd DGR
python main.py \
--dataset [nq, tqa, webq] \
--split [train, dev, test] \
--model_name [sentence-transformers/gtr-t5-large, sentence-transformers/all-MiniLM-L6-v2] \
--top_k [3, 5]
파라미터 설명
- --dataset: 사용할 데이터셋
- --split: 데이터 분할
- --model_name: 사용할 Sentence Transformer 모델
- --top_k: 검색할 상위 k개 문서 수
지원 모델
- sentence-transformers/gtr-t5-large: 고성능 T5 기반 모델
- sentence-transformers/all-MiniLM-L6-v2: 경량화된 모델
- 기타 Sentence Transformers 호환 모델
사용법
전체 파이프라인 실행
- 문서 생성 (DG 단계)
cd DG
python main.py --dataset nq --split train --engine text-davinci-003 --temperature 0.5 --max_tokens 300
- 문서 검색 (DGR 단계)
cd DGR
python main.py --dataset nq --split train --model_name sentence-transformers/gtr-t5-large --top_k 5
예제 실행
Natural Questions 데이터셋으로 실험:
# 1단계: 문서 생성
cd DG
python main.py --dataset nq --split dev --engine text-davinci-003 --num_sequence 10 --temperature 0.5 --max_tokens 200
# 2단계: 문서 검색
cd ../DGR
python main.py --dataset nq --split dev --model_name sentence-transformers/gtr-t5-large --top_k 3
데이터셋
이 프로젝트는 다음 오픈 도메인 질문 응답 데이터셋을 지원합니다:
- Natural Questions (NQ): Google 검색 쿼리 기반 질문-답변 데이터셋
- TriviaQA (TQA): 트리비아 질문과 Wikipedia/웹 문서 기반 답변
- WebQuestions (WEBQ): Freebase 기반 질문-답변 데이터셋
데이터 형식
{
"question": "질문 텍스트",
"answers": ["답변1", "답변2"],
"ctxs": [
{
"text": "문서 텍스트",
"score": 0.95
}
]
}
주요 의존성 패키지
- PyTorch: 딥러닝 프레임워크
- Transformers: Hugging Face 트랜스포머 라이브러리
- LlamaIndex: 문서 검색 및 인덱싱
- LangChain: LLM 애플리케이션 개발 프레임워크
- Sentence Transformers: 문장 임베딩
- OpenAI: GPT 모델 API 클라이언트
- Datasets: 데이터셋 처리
- NLTK: 자연어 처리
성능 및 결과
GRG 접근법은 기존의 검색 기반 방법들과 비교하여 다음과 같은 장점을 보입니다:
- 더 관련성 높은 문서 생성
- 향상된 검색 정확도
- 더 나은 최종 답변 품질
자세한 실험 결과는 논문을 참조하세요.
향후 계획
- Document Retriever (DR): 곧 공개 예정
- LLaMA Generator: 곧 공개 예정
문제 해결
일반적인 문제들
- OpenAI API 키 오류
- DG/utils.py에서 API 키가 올바르게 설정되었는지 확인하세요.
- 메모리 부족 오류
- --max_tokens 값을 줄이거나 배치 크기를 조정하세요.
- 모델 다운로드 오류
- 인터넷 연결을 확인하고 Hugging Face 모델에 대한 접근 권한을 확인하세요.
DPR
더보기
Dense Passage Retrieval (DPR)
Dense Passage Retrieval (DPR)은 최신 오픈 도메인 질의응답 연구를 위한 도구와 모델들의 집합입니다. 다음 논문을 기반으로 합니다:
모델 체크포인트를 기반으로 논문의 실험 결과를 재현하는 것에 관심이 있다면(즉, 인코더를 처음부터 훈련하고 싶지 않다면), pip를 통해 실험이 잘 패키징된 Pyserini 툴킷을 고려해보세요. 그들의 툴킷은 또한 더 높은 BM25 및 하이브리드 점수를 보고합니다.
주요 기능
- Dense retriever 모델은 bi-encoder 아키텍처를 기반으로 합니다.
- 이 논문에서 영감을 받은 추출형 Q&A reader & ranker 조인트 모델.
- 관련 데이터 전처리 및 후처리 도구.
- 추론 시간 로직을 위한 Dense retriever 컴포넌트는 FAISS 인덱스 기반입니다.
새로운 (2021년 3월) 릴리스
DPR 코드베이스가 여러 개선사항과 새로운 모델로 업그레이드되었습니다. 주요 변경사항:
- 데이터 로더를 제외한 모든 명령줄 도구에 대한 Hydra 기반 구성 (곧 변환 예정)
- 커스텀 데이터셋을 지원하는 플러그인 가능한 데이터 처리 레이어
- 더 나은 성능의 새로운 검색 모델 체크포인트.
새로운 (2021년 3월) 검색 모델
NQ 데이터셋에서만 훈련된 새로운 bi-encoder 모델이 제공됩니다:
새로운 체크포인트, 훈련 데이터, 검색 결과 및 위키피디아 임베딩. 이는 원본 DPR NQ 훈련 세트와 이전 NQ 체크포인트를 사용하여 DPR 인덱스 자체를 사용해 하드 네거티브가 마이닝된 버전에서 훈련되었습니다. Bi-encoder 모델은 이 새로운 훈련 데이터와 원본 NQ 훈련 데이터를 결합하여 처음부터 훈련됩니다. 이 훈련 방식은 좋은 검색 성능 향상을 제공합니다.
NQ 테스트 세트(3610개 질문)에서 새로운 모델 vs 기존 모델의 top-k 문서 검색 정확도.
| Top-k passages | 원본 DPR NQ 모델 | 새로운 DPR 모델 |
| 1 | 45.87 | 52.47 |
| 5 | 68.14 | 72.24 |
| 20 | 79.97 | 81.33 |
| 100 | 85.87 | 87.29 |
새 모델 다운로드 가능한 리소스 이름들 (아래 download_data 스크립트 사용법 참조):
체크포인트: checkpoint.retriever.single-adv-hn.nq.bert-base-encoder
새로운 훈련 데이터: data.retriever.nq-adv-hn-train
NQ 테스트 세트를 위한 검색기 결과: data.retriever_results.nq.single-adv-hn.test
위키피디아 임베딩: data.retriever_results.nq.single-adv-hn.wikipedia_passages
설치
소스에서 설치. Python의 가상환경이나 Conda 환경을 권장합니다.
git clone git@github.com:facebookresearch/DPR.git
cd DPR
pip install .
DPR은 Python 3.6+ 및 PyTorch 1.2.0+에서 테스트되었습니다. DPR은 인코더 코드 구현을 위해 서드파티 라이브러리에 의존합니다. 현재 Huggingface (버전 <=3.1.0) BERT, Pytext BERT 및 Fairseq RoBERTa 인코더 모델을 지원합니다. 토큰화 과정의 일반성으로 인해, DPR은 현재 Huggingface 토크나이저를 사용합니다. 따라서 Huggingface는 유일한 필수 종속성이며, Pytext & Fairseq는 선택사항입니다. 해당 인코더를 사용하려면 별도로 설치하세요.
리소스 및 데이터 형식
먼저, 검색기 또는 리더 훈련을 위한 데이터를 준비해야 합니다. 각 DPR 컴포넌트는 고유한 입출력 데이터 형식을 가지고 있습니다. 아래에서 형식 설명을 볼 수 있습니다. DPR은 dpr/data/download_data.py 도구를 사용하여 클라우드에서 다운로드할 수 있는 NQ & Trivia 전처리된 데이터셋(및 모델 체크포인트)을 제공합니다. 다운로드할 리소스 이름을 지정해야 합니다. 모든 옵션을 보려면 'python data/download_data.py'를 실행하세요.
python data/download_data.py \
--resource {download_data.py의 RESOURCES_MAP에서 키} \
[선택사항 --output_dir {당신의 위치}]
리소스 이름 매칭은 접두사 기반입니다. 모든 데이터 리소스를 다운로드해야 한다면, --resource data를 사용하면 됩니다.
검색기 입력 데이터 형식
검색기 훈련 데이터의 기본 데이터 형식은 JSON입니다. 질문당 2가지 유형의 네거티브 패시지 풀과 포지티브 패시지 및 일부 추가 정보를 포함합니다.
[
{
"question": "....",
"answers": ["...", "...", "..."],
"positive_ctxs": [{
"title": "...",
"text": "...."
}],
"negative_ctxs": ["..."],
"hard_negative_ctxs": ["..."]
},
...
]
negative_ctxs & hard_negative_ctxs의 요소 구조는 positive_ctxs와 정확히 동일합니다. 다운로드 가능한 전처리된 데이터는 모델 수정에 유용할 수 있는 일부 추가 속성(패시지당 bm25 점수 등)도 포함합니다. 하지만 현재 DPR에서는 사용되지 않습니다.
'data.retriever.nq' 키 접두사를 사용하여 논문에서 사용된 준비된 NQ 데이터셋을 다운로드할 수 있습니다. 이 형식에서는 dev & train 서브셋만 사용 가능합니다. 또한 모든 train/dev/test 분할에 대한 질문 & 답변만 있는 CSV 데이터 파일도 제공합니다. 이들은 NQ 전처리 단계에서 원본 샘플 세트의 일부가 손실되므로 모델 평가에 사용됩니다. 평가를 위한 각 세트를 얻으려면 'data.retriever.qas.*' 리소스 키를 사용하세요.
python data/download_data.py
--resource data.retriever
[선택사항 --output_dir {당신의 위치}]
DPR 데이터 형식 및 사용자 정의 처리
dpr/data/{biencoder|retriever|reader}_data.py 파일의 DPR Dataset 클래스를 상속하고 load_data() 및 getitem() 메서드를 구현하여 자신만의 데이터 형식과 사용자 정의 데이터 파싱 & 로딩 로직을 사용할 수 있습니다. DPR hydra 구성 지침을 참조하세요.
검색기 훈련
검색기 훈련 품질은 효과적인 배치 크기에 따라 달라집니다. 논문에서 보고된 것은 8 x 32GB GPU를 사용했습니다. 한 머신에서 훈련을 시작하려면:
python train_dense_encoder.py \
train_datasets=[훈련 데이터셋 목록, 공백 없이 쉼표로 구분] \
dev_datasets=[개발 데이터셋 목록, 공백 없이 쉼표로 구분] \
train=biencoder_local \
output_dir={체크포인트 디렉터리 경로}
NQ 데이터셋 예시
python train_dense_encoder.py \
train_datasets=[nq_train] \
dev_datasets=[nq_dev] \
train=biencoder_local \
output_dir={체크포인트 디렉터리 경로}
DPR은 기본적으로 HuggingFace BERT-base를 인코더로 사용합니다. 다른 준비된 옵션으로는 Fairseq의 ROBERTA 및 Pytext BERT 모델이 있습니다. 인코더 구성 파일(conf/encoder/hf_bert.yaml)을 변경하거나 conf/encoder 디렉터리에 새 구성 파일을 제공하고 encoder={새 파일 이름} 명령줄 매개변수로 활성화하여 선택할 수 있습니다.
참고사항:
- pytext bert 또는 fairseq roberta를 사용하려면 사전 훈련된 가중치를 다운로드하고 encoder.pretrained_file 매개변수를 지정해야 합니다. RoBERTa 모델의 경우 'pretrained.fairseq.roberta-base' 리소스 접두사에 대한 다운로드된 파일의 디렉터리 위치를 지정하거나 pytext BERT의 경우 파일 경로를 지정하세요(리소스 이름 'pretrained.pytext.bert-base.model').
- 검증 및 체크포인트 저장은 train.eval_per_epoch 매개변수 값에 따라 발생합니다.
- 지정된 에포크 수(train.num_train_epochs 구성 매개변수) 외에는 중지 조건이 없습니다.
- 모든 평가는 모델 체크포인트를 저장합니다.
- 최고 체크포인트는 훈련 과정 출력에 로그됩니다.
- bi-encoder 훈련을 위한 정규 NLL 분류 손실 검증은 평균 순위 평가로 대체될 수 있습니다. 이는 입력 데이터 패시지 풀에서 패시지와 질문 벡터를 집계하고, 이러한 표현에 대해 큰 유사성 행렬 계산을 수행한 다음 각 질문에 대한 골드 패시지의 순위를 평균화합니다. 이 메트릭이 nll 분류 손실과 비교하여 최종 검색 성능과 더 상관관계가 있음을 발견했습니다. 그러나 이 평균 순위 검증은 DistributedDataParallel vs DataParallel PyTorch 모드에서 다르게 작동합니다. 이 모드를 활성화하고 설정을 수정하려면 train.val_av_rank_* 매개변수 세트를 참조하세요.
최고 하이퍼파라미터 설정에 대한 e2e 예시는 아래 섹션을 참조하세요.
검색기 추론
정적 문서 데이터셋에 대한 표현 벡터 생성은 단일 GPU에서 계산할 경우 며칠이 걸릴 수 있는 매우 병렬화 가능한 프로세스입니다. 각각 독립적으로 스크립트를 실행하고 자신만의 샤드를 지정하여 여러 사용 가능한 GPU 서버를 사용할 수 있습니다.
python generate_dense_embeddings.py \
model_file={biencoder 체크포인트 경로} \
ctx_src={패시지 리소스 이름, 원본 위키피디아 분할을 사용하려면 dpr_wiki로 설정} \
shard_id={샤드 번호, 0부터 시작} num_shards={총 샤드 수} \
out_file={결과 파일 위치 + 이름 접두사}
ctx_src 매개변수의 리소스 이름은 conf/ctx_sources/default_sources.yaml 파일의 소스 이름입니다.
참고: 여기서는 훈련 모드에 비해 훨씬 큰 배치 크기를 사용할 수 있습니다. 예를 들어, 2 GPU(16gb) 서버에 대해 batch_size 128을 설정하면 잘 작동합니다. 원본 모델(NQ 데이터셋에서 훈련)에서 이미 생성된 위키피디아 임베딩을 리소스 키 'data.retriever_results.nq.single.wikipedia_passages'를 사용하여 다운로드할 수 있습니다. 새로운 더 나은 모델의 임베딩 리소스 이름 'data.retriever_results.nq.single-adv-hn.wikipedia_passages'
일반적으로 50개의 2-gpu 노드에서 다음 매개변수를 사용합니다: batch_size=128 shard_id=0 num_shards=50
전체 문서 세트에 대한 검색기 검증:
python dense_retriever.py \
model_file={download_data.py에서 'checkpoint.retriever.single.nq.bert-base-encoder'로 다운로드된 체크포인트 경로} \
qa_dataset={테스트 소스 이름} \
ctx_datatsets=[{패시지 소스 이름 목록, 공백 없이 쉼표로 구분}] \
encoded_ctx_files=[{인코딩된 문서 파일 glob 표현식 목록, 공백 없이 쉼표로 구분}] \
out_file={결과가 있는 출력 json 파일 경로}
예를 들어, 두 패시지 세트에 대해 생성된 임베딩이 ~/myproject/embeddings_passages1/wiki_passages_* 및 ~/myproject/embeddings_passages2/wiki_passages_* 파일이고 NQ 데이터셋에서 평가하려면:
python dense_retriever.py \
model_file={체크포인트 파일 경로} \
qa_dataset=nq_test \
ctx_datatsets=[dpr_wiki] \
encoded_ctx_files=[\"~/myproject/embeddings_passages1/wiki_passages_*\",\"~/myproject/embeddings_passages2/wiki_passages_*\"] \
out_file={출력 json 파일 경로}
이 도구는 후속 리더 모델 훈련을 위해 검색된 결과를 지정된 out_file에 작성합니다. 다음 형식의 json입니다:
[
{
"question": "...",
"answers": ["...", "...", ... ],
"ctxs": [
{
"id": "...", # 데이터베이스 tsv 파일의 패시지 id
"title": "",
"text": "....",
"score": "...", # 검색기 점수
"has_answer": true|false
},
]
결과는 유사성 점수에 따라 가장 관련성이 높은 것부터 가장 낮은 것 순으로 정렬됩니다.
기본적으로 dense_retriever는 전수 검색 프로세스를 사용하지만, 손실 인덱스 유형을 사용하도록 선택할 수 있습니다. HNSW 및 HNSW_SQ 인덱스 옵션을 제공합니다. indexer=hnsw 또는 indexer=hnsw_sq 명령줄 인수로 활성화하세요. 이 인덱스를 사용하는 것은 빠른 검색 프로세스가 훨씬 긴 인덱싱 시간과 높은 RAM 사용량의 비용으로 제공되므로 연구 관점에서는 무용할 수 있습니다. 제공되는 유사성 점수는 전수 검색의 기본 경우(indexer=flat)에는 내적이고 HNSW 인덱스의 경우 수정된 표현 공간에서의 L2 거리입니다.
리더 모델 훈련
python train_extractive_reader.py \
encoder.sequence_length=350 \
train_files={검색기 훈련 세트 결과 파일 경로} \
dev_files={검색기 개발 세트 결과 파일 경로} \
output_dir={출력 디렉터리 경로}
기본 하이퍼파라미터는 8 gpu가 있는 단일 노드 설정에 맞춰져 있습니다. conf/train/extractive_reader_default.yaml 및 conf/extractive_reader_train_cfg.yaml 구성 파일에서 필요에 따라 수정하거나 명령줄에서 특정 매개변수를 재정의하세요. 첫 번째 실행은 train_files & dev_files를 전처리하고 동일한 위치에서 직렬화된 .pkl 파일 세트로 변환하며 모든 후속 실행에서 이를 사용합니다.
참고사항:
- pytext bert 또는 fairseq roberta를 사용하려면 사전 훈련된 가중치를 다운로드하고 encoder.pretrained_file 매개변수를 지정해야 합니다. RoBERTa 모델의 경우 'pretrained.fairseq.roberta-base' 리소스 접두사에 대한 다운로드된 파일의 디렉터리 위치를 지정하거나 pytext BERT의 경우 파일 경로를 지정하세요(리소스 이름 'pretrained.pytext.bert-base.model').
- 리더 훈련 파이프라인은 train.eval_step 배치마다 모델 검증을 수행합니다
- bi-encoder와 마찬가지로 모든 검증에서 모델 체크포인트를 저장합니다
- bi-encoder와 마찬가지로 지정된 에포크 수 외에는 중지 조건이 없습니다.
- bi-encoder와 마찬가지로 최고 체크포인트 선택 로직이 없으므로 훈련 과정 출력에 로그된 개발 세트 검증 성능을 기반으로 선택해야 합니다.
- 현재 코드는 Exact Match 메트릭만 계산합니다.
리더 모델 추론
추론을 수행하려면 train_files를 지정하지 않고 train_reader.py를 실행하세요. 체크포인트 경로가 있는 model_file, 질문당 패시지 수가 있는 passages_per_question_predict(예측 파일을 저장할 때 사용됨), 질문의 답변 스팬을 선택할 top 패시지 임계값 목록이 있는 eval_top_docs(로그로 출력됨)를 지정해야 합니다. 명령줄 예시는 다음과 같습니다.
python train_extractive_reader.py \
prediction_results_file={결과를 작성할 파일 경로} \
eval_top_docs=[10,20,40,50,80,100] \
dev_files={평가할 검색기 결과 파일 경로} \
model_file= {리더 체크포인트 경로} \
train.dev_batch_size=80 \
passages_per_question_predict=100 \
encoder.sequence_length=350
분산 훈련
Pytorch의 분산 훈련 런처 도구를 사용하세요:
python -m torch.distributed.launch \
--nproc_per_node={WORLD_SIZE} {비분산 스크립트 이름 & 매개변수}
참고:
- 모든 배치 크기 관련 매개변수는 분산 모드(DistributedDataParallel)에서는 gpu당으로 지정되고 DataParallel(단일 노드 - 다중 gpu) 모드에서는 모든 사용 가능한 gpu에 대해 지정됩니다.
최고 하이퍼파라미터 설정
NQ 데이터셋에 대한 최고 설정의 e2e 예시.
1. 모든 검색기 훈련 및 검증 데이터 다운로드:
python data/download_data.py --resource data.wikipedia_split.psgs_w100
python data/download_data.py --resource data.retriever.nq
python data/download_data.py --resource data.retriever.qas.nq
2. 단일 세트 모드에서 Biencoder(검색기) 훈련.
단일 8 GPU x 32 GB 서버에서 분산 훈련 모드를 사용했습니다
python -m torch.distributed.launch --nproc_per_node=8
train_dense_encoder.py \
train=biencoder_nq \
train_datasets=[nq_train] \
dev_datasets=[nq_dev] \
train=biencoder_nq \
output_dir={당신의 출력 디렉터리}
새 모델 훈련은 두 NQ 데이터셋을 결합합니다:
python -m torch.distributed.launch --nproc_per_node=8
train_dense_encoder.py \
train=biencoder_nq \
train_datasets=[nq_train,nq_train_hn1] \
dev_datasets=[nq_dev] \
train=biencoder_nq \
output_dir={당신의 출력 디렉터리}
40 에포크 훈련을 완료하는 데 약 하루가 걸립니다. 에포크 30에서 평균 순위 검증으로 전환하며 끝에서 25 이하여야 합니다. bi-encoder의 최고 체크포인트는 보통 마지막이지만, 에포크 ~25 이후의 것을 가져와도 그리 다르지 않을 것입니다.
3. 위키피디아에 대한 임베딩 생성.
"대용량 문서 세트에 대한 표현 생성" 지침을 사용하세요. 50개의 2 GPU 서버에서 2천 1백만 패시지 표현 벡터를 생성하는 데 약 40분이 걸립니다.
4. 검색 정확도를 평가하고 각 train/dev/test 데이터셋에 대한 top 패시지 결과를 생성합니다.
python dense_retriever.py \
model_file={최고 체크포인트 경로 또는 제공된 체크포인트 사용 (checkpoint.retriever.*와 같은 리소스 이름)} \
qa_dataset=nq_test \
ctx_datatsets=[dpr_wiki] \
encoded_ctx_files=["{생성된 임베딩 파일에 대한 glob 표현식}"] \
out_file={출력 파일 경로}
사용 가능한 GPU 수에 따라 batch_size를 조정하세요. 2 GPU 서버의 경우 64-128이 작동할 것입니다.
5. 리더 훈련
단일 8 GPU x 32 GB 서버를 사용하여 대용량 데이터셋에 대한 리더 모델을 훈련했습니다. 모든 기본 매개변수는 이미 최고 NQ 설정으로 설정되어 있습니다. NQ 데이터셋에 대한 data.gold_passages_info.nq_train & data.gold_passages_info.nq_dev 리소스도 다운로드하세요 - 이들은 NQ 리더 훈련을 위해 데이터를 전처리할 때 NQ 전용 휴리스틱에 사용됩니다. gold_passages_src & gold_passages_src_dev가 지정되지 않은 상태에서 NQ 데이터에 대해 리더 훈련을 이미 실행했다면, 해당 .pkl 파일을 삭제하여 재생성되도록 하세요.
python train_extractive_reader.py \
encoder.sequence_length=350 \
train_files={검색기 훈련 세트 결과 파일 경로} \
dev_files={검색기 개발 세트 결과 파일 경로} \
gold_passages_src={data.gold_passages_info.nq_train 파일 경로} \
gold_passages_src_dev={data.gold_passages_info.nq_dev 파일 경로} \
output_dir={출력 디렉터리 경로}
위의 학습률을 정적 스케줄과 함께 사용하는 것이 가장 효과적임을 발견했으므로 평가 성능 역학을 기반으로 수동으로 훈련을 중지해야 합니다. 최고 결과는 16-18 훈련 에포크 또는 약 6만 모델 업데이트 후에 달성되었습니다.
NQ 데이터셋에 대한 e2e 파이프라인의 모든 입력 및 중간 결과와 Trivia에 대한 대부분의 유사한 리소스를 제공합니다.
기타
- TREC 검증은 정규표현식 기반 매칭이 필요합니다. 정규표현식 모드에서는 검색기 검증만 지원합니다. --match 매개변수 옵션을 참조하세요.
- WebQ 검증은 현재 포함되지 않은 엔티티 정규화가 필요합니다.
데이터셋
Natural Questions (NQ)
- 설명: Google에서 수집한 실제 검색 질의 기반 질의응답 데이터셋
- 사용: 주요 훈련 및 평가 데이터셋으로 사용
- 형식: JSON 형식으로 질문, 답변, 관련 패시지 포함
TriviaQA
- 설명: 퀴즈 질문 기반 읽기 이해 데이터셋
- 사용: 추가적인 평가 및 훈련 데이터로 활용
Wikipedia
- 설명: 위키피디아 문서들을 패시지 단위로 분할한 지식 베이스
- 사용: 검색 대상 문서 집합으로 활용
- 처리: 100단어 단위로 분할된 패시지들
모델 아키텍처
Bi-Encoder (Dense Retriever)
- 구조: 질문과 패시지를 각각 독립적으로 인코딩하는 이중 인코더
- 기반 모델: BERT-base, RoBERTa, Pytext BERT 지원
- 목적: 질문과 패시지 간의 의미적 유사도 계산
Extractive Reader
- 구조: 검색된 패시지에서 정확한 답변 스팬을 추출하는 모델
- 기능: 패시지 랭킹과 답변 추출을 동시에 수행
- 기반: BERT 기반 모델 사용
전체 파이프라인 흐름
- 데이터 전처리: NQ, TriviaQA 등의 원시 데이터를 DPR 형식으로 변환
- Dense Retriever 훈련: 질문-패시지 쌍으로 bi-encoder 모델 훈련
- 임베딩 생성: 전체 위키피디아 패시지에 대한 dense vector 생성
- 검색 성능 평가: 테스트 세트에서 top-k 검색 정확도 측정
- Reader 훈련: 검색된 패시지에서 답변 추출을 위한 모델 훈련
- End-to-End 평가: 검색과 읽기를 결합한 최종 성능 평가
주요 개선사항 (2021년 3월 업데이트)
- Hard Negative Mining: DPR 인덱스 자체를 사용한 어려운 네거티브 샘플 마이닝
- 향상된 성능: 새로운 모델에서 기존 대비 6-7% 성능 향상
- Hydra 통합: 구성 관리를 위한 Hydra 프레임워크 도입
- 플러그인 데이터 처리: 커스텀 데이터셋 지원을 위한 모듈화된 구조
https://www.kaggle.com/code/leireher/rag-for-triviaqa
RAG for TriviaQA
Explore and run machine learning code with Kaggle Notebooks | Using data from multiple data sources
www.kaggle.com
"""
Wikipedia RAG (Retrieval-Augmented Generation) System for TriviaQA
이 시스템은 영어 Wikipedia 데이터를 수집하여 청킹, 임베딩하고
TriviaQA 질문-답변 시스템을 위한 컨텍스트를 제공합니다.
주요 기능:
1. Wikipedia 데이터 수집 및 전처리
2. 텍스트 청킹 및 임베딩 생성
3. FAISS 및 ChromaDB 벡터 검색
4. TriviaQA 데이터셋 연동
"""
import os
import re
import pickle
from typing import List, Dict, Optional, Any
from dataclasses import dataclass
import numpy as np
import pandas as pd
import wikipediaapi
import nltk
import faiss
import chromadb
from tqdm import tqdm
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
from datasets import load_dataset
# =====================================================================
# 설정 및 데이터 클래스
# =====================================================================
@dataclass
class WikipediaRAGConfig:
"""Wikipedia RAG 시스템 설정"""
# 모델 설정
embedding_model_name: str = "all-MiniLM-L6-v2"
# 청킹 설정
chunk_size: int = 512
chunk_overlap: int = 50
# Wikipedia 설정
max_chars_per_page: int = 50000
user_agent: str = "RAG-Dataset-Creator/1.0 (https://example.com/contact)"
# 저장 경로
pickle_file: str = "wikipedia_rag_dataset.pkl"
faiss_index_file: str = "wikipedia_faiss.index"
chromadb_path: str = "./chroma_db"
chromadb_collection: str = "wikipedia_collection"
# 임베딩 설정
embedding_batch_size: int = 32
search_top_k: int = 5
@dataclass
class SearchResult:
"""검색 결과 데이터 클래스"""
rank: int
score: float
document: str
metadata: Dict[str, Any]
@dataclass
class WikipediaContent:
"""Wikipedia 페이지 내용 데이터 클래스"""
title: str
content: str
url: str
summary: str
@dataclass
class DocumentChunk:
"""문서 청크 데이터 클래스"""
text: str
metadata: Dict[str, Any]
# =====================================================================
# 유틸리티 함수
# =====================================================================
def setup_nltk():
"""NLTK 데이터 다운로드"""
try:
nltk.download('punkt', quiet=True)
nltk.download('stopwords', quiet=True)
print("✅ NLTK 데이터 다운로드 완료")
except Exception as e:
print(f"⚠️ NLTK 데이터 다운로드 중 오류: {e}")
def get_default_wikipedia_titles() -> List[str]:
"""기본 Wikipedia 페이지 제목 리스트 반환"""
return [
"Artificial Intelligence",
"Machine Learning",
"Natural Language Processing",
"Computer Science",
"Physics",
"Mathematics",
"Biology",
"Chemistry",
"History",
"Geography",
"Literature",
"Philosophy",
"Psychology",
"Economics",
"Politics",
"Sports",
"Music",
"Art",
"Science",
"Technology",
"World War II",
"Ancient Egypt",
"Roman Empire",
"Renaissance",
"Industrial Revolution",
"Space exploration",
"Climate change",
"Evolution",
"Quantum mechanics",
"Einstein"
]
def get_test_wikipedia_titles() -> List[str]:
"""테스트용 Wikipedia 페이지 제목 리스트 반환"""
return [
"Artificial Intelligence",
"Machine Learning",
"Computer Science",
"Physics",
"Mathematics"
]
# =====================================================================
# 메인 Wikipedia RAG 데이터셋 클래스
# =====================================================================
class WikipediaRAGDataset:
"""Wikipedia RAG 데이터셋 생성 및 관리 클래스"""
def __init__(self, config: WikipediaRAGConfig = None):
"""
초기화
Args:
config: 설정 객체 (None이면 기본 설정 사용)
"""
self.config = config or WikipediaRAGConfig()
# Wikipedia API 초기화
self.wiki_api = wikipediaapi.Wikipedia(
language='en',
extract_format=wikipediaapi.ExtractFormat.WIKI,
user_agent=self.config.user_agent
)
# 임베딩 모델 초기화
self.embedding_model = SentenceTransformer(self.config.embedding_model_name)
# 텍스트 분할기 초기화
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=self.config.chunk_size,
chunk_overlap=self.config.chunk_overlap,
length_function=len,
separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
)
# 데이터 저장 변수
self.documents: List[str] = []
self.embeddings: Optional[np.ndarray] = None
self.metadata: List[Dict[str, Any]] = []
print("✅ WikipediaRAGDataset 초기화 완료!")
def get_wikipedia_content(self, titles: List[str]) -> List[WikipediaContent]:
"""
Wikipedia 페이지들의 내용을 수집합니다.
Args:
titles: Wikipedia 페이지 제목 리스트
Returns:
WikipediaContent 객체 리스트
"""
contents = []
for title in tqdm(titles, desc="Wikipedia 페이지 수집"):
try:
page = self.wiki_api.page(title)
if page.exists():
content = page.text
# 페이지 크기 제한
if len(content) > self.config.max_chars_per_page:
content = content[:self.config.max_chars_per_page]
wiki_content = WikipediaContent(
title=title,
content=content,
url=page.fullurl,
summary=page.summary
)
contents.append(wiki_content)
print(f"✅ '{title}' 페이지 수집 완료 (길이: {len(content)})")
else:
print(f"❌ '{title}' 페이지를 찾을 수 없습니다.")
except Exception as e:
print(f"❌ '{title}' 페이지 수집 중 오류: {e}")
return contents
def chunk_documents(self, contents: List[WikipediaContent]) -> List[DocumentChunk]:
"""
문서들을 청킹합니다.
Args:
contents: WikipediaContent 객체 리스트
Returns:
DocumentChunk 객체 리스트
"""
all_chunks = []
for content in tqdm(contents, desc="문서 청킹"):
# 텍스트 청킹
text_chunks = self.text_splitter.split_text(content.content)
for i, chunk_text in enumerate(text_chunks):
chunk_metadata = {
'source': content.title,
'chunk_id': i,
'url': content.url,
'summary': content.summary
}
chunk = DocumentChunk(
text=chunk_text,
metadata=chunk_metadata
)
all_chunks.append(chunk)
print(f"✅ 총 {len(all_chunks)}개의 청크 생성 완료!")
return all_chunks
def create_embeddings(self, chunks: List[DocumentChunk]) -> np.ndarray:
"""
청크들의 임베딩을 생성합니다.
Args:
chunks: DocumentChunk 객체 리스트
Returns:
임베딩 배열
"""
texts = [chunk.text for chunk in chunks]
print("🔄 임베딩 생성 중...")
embeddings = self.embedding_model.encode(
texts,
batch_size=self.config.embedding_batch_size,
show_progress_bar=True,
convert_to_numpy=True
)
# 데이터 저장
self.documents = [chunk.text for chunk in chunks]
self.embeddings = embeddings
self.metadata = [chunk.metadata for chunk in chunks]
print(f"✅ 임베딩 생성 완료! 형태: {embeddings.shape}")
return embeddings
def save_to_pickle(self, filename: str = None):
"""
데이터를 pickle 파일로 저장합니다.
Args:
filename: 저장할 파일명 (None이면 설정값 사용)
"""
filename = filename or self.config.pickle_file
data = {
'documents': self.documents,
'embeddings': self.embeddings,
'metadata': self.metadata,
'embedding_dim': self.embedding_model.get_sentence_embedding_dimension(),
'config': self.config
}
with open(filename, 'wb') as f:
pickle.dump(data, f)
print(f"✅ 데이터가 '{filename}'에 저장되었습니다.")
def create_faiss_index(self, index_file: str = None) -> faiss.Index:
"""
FAISS 인덱스를 생성합니다.
Args:
index_file: 인덱스 파일명 (None이면 설정값 사용)
Returns:
FAISS 인덱스
"""
index_file = index_file or self.config.faiss_index_file
if self.embeddings is None:
raise ValueError("임베딩이 생성되지 않았습니다. create_embeddings()를 먼저 실행하세요.")
# FAISS 인덱스 생성
dimension = self.embeddings.shape[1]
index = faiss.IndexFlatIP(dimension) # Inner Product (cosine similarity)
# 임베딩 정규화 (cosine similarity를 위해)
embeddings_normalized = self.embeddings / np.linalg.norm(
self.embeddings, axis=1, keepdims=True
)
# 인덱스에 임베딩 추가
index.add(embeddings_normalized.astype('float32'))
# 인덱스 저장
faiss.write_index(index, index_file)
print(f"✅ FAISS 인덱스가 '{index_file}'에 저장되었습니다.")
return index
def save_to_chromadb(self, collection_name: str = None,
persist_directory: str = None) -> chromadb.Collection:
"""
ChromaDB에 데이터를 저장합니다.
Args:
collection_name: 컬렉션명 (None이면 설정값 사용)
persist_directory: 저장 디렉토리 (None이면 설정값 사용)
Returns:
ChromaDB 컬렉션
"""
collection_name = collection_name or self.config.chromadb_collection
persist_directory = persist_directory or self.config.chromadb_path
if not self.documents:
raise ValueError("문서가 없습니다. 먼저 데이터를 생성하세요.")
# ChromaDB 클라이언트 생성
client = chromadb.PersistentClient(path=persist_directory)
# 기존 컬렉션 삭제 (있다면)
try:
client.delete_collection(collection_name)
except:
pass
# 새 컬렉션 생성
collection = client.create_collection(collection_name)
# 데이터 추가
ids = [f"doc_{i}" for i in range(len(self.documents))]
collection.add(
documents=self.documents,
embeddings=self.embeddings.tolist(),
metadatas=self.metadata,
ids=ids
)
print(f"✅ ChromaDB에 {len(self.documents)}개의 문서가 저장되었습니다.")
return collection
def create_full_dataset(self, titles: List[str] = None) -> 'WikipediaRAGDataset':
"""
전체 데이터셋 생성 파이프라인을 실행합니다.
Args:
titles: Wikipedia 페이지 제목 리스트 (None이면 기본값 사용)
Returns:
자기 자신 (method chaining을 위해)
"""
titles = titles or get_default_wikipedia_titles()
print(f"📚 총 {len(titles)}개의 Wikipedia 페이지로 데이터셋 생성 시작!")
# 1. Wikipedia 내용 수집
print("\n=== 1단계: Wikipedia 내용 수집 ===")
contents = self.get_wikipedia_content(titles)
# 2. 문서 청킹
print("\n=== 2단계: 문서 청킹 ===")
chunks = self.chunk_documents(contents)
# 3. 임베딩 생성
print("\n=== 3단계: 임베딩 생성 ===")
embeddings = self.create_embeddings(chunks)
# 4. 데이터 저장
print("\n=== 4단계: 데이터 저장 ===")
self.save_to_pickle()
self.create_faiss_index()
self.save_to_chromadb()
print("\n🎉 데이터셋 생성 완료!")
print(f"- 총 문서 수: {len(self.documents)}")
print(f"- 임베딩 차원: {embeddings.shape[1]}")
return self
# =====================================================================
# Wikipedia RAG 검색 클래스
# =====================================================================
class WikipediaRAGRetriever:
"""Wikipedia RAG 검색 클래스"""
def __init__(self, config: WikipediaRAGConfig = None,
pickle_file: str = None, faiss_index_file: str = None):
"""
초기화
Args:
config: 설정 객체
pickle_file: 데이터셋 파일 경로
faiss_index_file: FAISS 인덱스 파일 경로
"""
self.config = config or WikipediaRAGConfig()
pickle_file = pickle_file or self.config.pickle_file
faiss_index_file = faiss_index_file or self.config.faiss_index_file
# 데이터 로드
print("📂 데이터셋 로드 중...")
self._load_dataset(pickle_file)
# FAISS 인덱스 로드
print("🔍 FAISS 인덱스 로드 중...")
self.faiss_index = faiss.read_index(faiss_index_file)
# 임베딩 모델 초기화
self.embedding_model = SentenceTransformer(self.config.embedding_model_name)
print(f"✅ 데이터셋 로드 완료! 총 {len(self.documents)}개의 문서")
def _load_dataset(self, pickle_file: str):
"""데이터셋 파일 로드"""
with open(pickle_file, 'rb') as f:
data = pickle.load(f)
self.documents = data['documents']
self.embeddings = data['embeddings']
self.metadata = data['metadata']
# 설정 정보가 있다면 업데이트
if 'config' in data:
self.config = data['config']
def search_similar_documents(self, query: str, top_k: int = None) -> List[SearchResult]:
"""
쿼리와 유사한 문서들을 검색합니다.
Args:
query: 검색 쿼리
top_k: 반환할 상위 결과 수
Returns:
SearchResult 객체 리스트
"""
top_k = top_k or self.config.search_top_k
# 쿼리 임베딩 생성
query_embedding = self.embedding_model.encode([query])
query_embedding = query_embedding / np.linalg.norm(
query_embedding, axis=1, keepdims=True
)
# FAISS로 유사한 문서 검색
scores, indices = self.faiss_index.search(
query_embedding.astype('float32'), top_k
)
# 결과 생성
results = []
for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
result = SearchResult(
rank=i + 1,
score=float(score),
document=self.documents[idx],
metadata=self.metadata[idx]
)
results.append(result)
return results
def search_with_chromadb(self, query: str, top_k: int = None,
collection_name: str = None,
persist_directory: str = None) -> List[SearchResult]:
"""
ChromaDB를 사용하여 검색합니다.
Args:
query: 검색 쿼리
top_k: 반환할 상위 결과 수
collection_name: 컬렉션명
persist_directory: 데이터베이스 디렉토리
Returns:
SearchResult 객체 리스트
"""
top_k = top_k or self.config.search_top_k
collection_name = collection_name or self.config.chromadb_collection
persist_directory = persist_directory or self.config.chromadb_path
# ChromaDB 클라이언트 생성
client = chromadb.PersistentClient(path=persist_directory)
collection = client.get_collection(collection_name)
# 검색 실행
results = collection.query(
query_texts=[query],
n_results=top_k
)
# 결과 생성
search_results = []
for i in range(len(results['documents'][0])):
result = SearchResult(
rank=i + 1,
score=results['distances'][0][i],
document=results['documents'][0][i],
metadata=results['metadatas'][0][i]
)
search_results.append(result)
return search_results
def pretty_print_results(self, results: List[SearchResult], query: str):
"""
검색 결과를 보기 좋게 출력합니다.
Args:
results: 검색 결과 리스트
query: 검색 쿼리
"""
print(f"\n🔍 검색 쿼리: '{query}'")
print("=" * 80)
for result in results:
print(f"\n📄 순위 {result.rank} (점수: {result.score:.4f})")
print(f"📂 출처: {result.metadata['source']}")
print(f"📝 내용: {result.document[:300]}...")
if len(result.document) > 300:
print(" [더 많은 내용...]")
print("-" * 80)
# =====================================================================
# TriviaQA 연동 클래스
# =====================================================================
class TriviaQAWithWikipediaRAG:
"""TriviaQA와 Wikipedia RAG 연동 클래스"""
def __init__(self, retriever: WikipediaRAGRetriever):
"""
초기화
Args:
retriever: Wikipedia RAG 검색기
"""
self.retriever = retriever
def get_context_for_question(self, question: str, top_k: int = None) -> Dict[str, Any]:
"""
질문에 대한 관련 컨텍스트를 Wikipedia에서 검색합니다.
Args:
question: 질문
top_k: 검색할 상위 결과 수
Returns:
컨텍스트 정보가 포함된 딕셔너리
"""
top_k = top_k or self.retriever.config.search_top_k
# Wikipedia에서 관련 문서 검색
search_results = self.retriever.search_similar_documents(question, top_k)
# 컨텍스트 생성
context_parts = []
for result in search_results:
context_parts.append(f"[{result.metadata['source']}] {result.document}")
context = "\n\n".join(context_parts)
return {
'question': question,
'context': context,
'search_results': search_results
}
def process_triviaqa_sample(self, sample: Dict[str, Any], top_k: int = None) -> Dict[str, Any]:
"""
TriviaQA 샘플을 처리하여 Wikipedia 컨텍스트를 추가합니다.
Args:
sample: TriviaQA 데이터 샘플
top_k: 검색할 상위 결과 수
Returns:
처리된 샘플 데이터
"""
top_k = top_k or self.retriever.config.search_top_k
# 질문과 답변 추출
question = sample.get('question', '')
answer = None
# 답변 추출 (다양한 형태 처리)
if 'answer' in sample:
answer_data = sample['answer']
if isinstance(answer_data, dict) and 'value' in answer_data:
answer = answer_data['value']
elif isinstance(answer_data, str):
answer = answer_data
# Wikipedia 컨텍스트 검색
context_data = self.get_context_for_question(question, top_k)
return {
'question': question,
'answer': answer,
'wikipedia_context': context_data['context'],
'search_results': context_data['search_results'],
'original_sample': sample
}
# =====================================================================
# 메인 실행 함수들
# =====================================================================
def create_wikipedia_rag_dataset(config: WikipediaRAGConfig = None,
titles: List[str] = None,
is_test: bool = False) -> WikipediaRAGDataset:
"""
Wikipedia RAG 데이터셋을 생성합니다.
Args:
config: 설정 객체
titles: Wikipedia 페이지 제목 리스트
is_test: 테스트 모드 여부
Returns:
생성된 데이터셋 객체
"""
config = config or WikipediaRAGConfig()
if is_test:
titles = titles or get_test_wikipedia_titles()
config.chunk_size = 256
config.chunk_overlap = 25
config.pickle_file = "test_wikipedia_rag.pkl"
config.faiss_index_file = "test_wikipedia_faiss.index"
else:
titles = titles or get_default_wikipedia_titles()
# 데이터셋 생성
dataset = WikipediaRAGDataset(config)
dataset.create_full_dataset(titles)
return dataset
def test_wikipedia_search(config: WikipediaRAGConfig = None,
pickle_file: str = None,
faiss_index_file: str = None):
"""
Wikipedia RAG 검색 기능을 테스트합니다.
Args:
config: 설정 객체
pickle_file: 데이터셋 파일 경로
faiss_index_file: FAISS 인덱스 파일 경로
"""
try:
# 검색기 초기화
retriever = WikipediaRAGRetriever(config, pickle_file, faiss_index_file)
# 테스트 쿼리들
test_queries = [
"artificial intelligence and machine learning",
"quantum mechanics and physics",
"World War II history",
"ancient civilizations",
"climate change and environment"
]
print("🔍 Wikipedia RAG 검색 테스트 시작!")
for query in test_queries:
print(f"\n{'='*100}")
# FAISS 검색 테스트
results = retriever.search_similar_documents(query, top_k=3)
retriever.pretty_print_results(results, query)
# ChromaDB 검색 테스트 (파일이 있는 경우)
try:
chroma_results = retriever.search_with_chromadb(query, top_k=3)
print(f"\n🎯 ChromaDB 검색 결과:")
retriever.pretty_print_results(chroma_results, query)
except Exception as e:
print(f"ChromaDB 검색 건너뜀: {e}")
print("\n✅ 검색 테스트 완료!")
except Exception as e:
print(f"❌ 검색 테스트 중 오류 발생: {e}")
print("먼저 데이터셋을 생성해주세요.")
def demonstrate_triviaqa_wikipedia_integration(config: WikipediaRAGConfig = None,
pickle_file: str = None,
faiss_index_file: str = None):
"""
TriviaQA와 Wikipedia RAG 연동을 시연합니다.
Args:
config: 설정 객체
pickle_file: 데이터셋 파일 경로
faiss_index_file: FAISS 인덱스 파일 경로
"""
print("🎯 TriviaQA + Wikipedia RAG 연동 시연")
print("=" * 60)
try:
# TriviaQA 데이터셋 로드
print("📚 TriviaQA 데이터셋 로드 중...")
ds = load_dataset("mandarjoshi/trivia_qa", "rc")
# Wikipedia RAG 검색기 초기화
print("🔍 Wikipedia RAG 검색기 초기화 중...")
retriever = WikipediaRAGRetriever(config, pickle_file, faiss_index_file)
# TriviaQA + Wikipedia RAG 연동기 초기화
trivia_rag = TriviaQAWithWikipediaRAG(retriever)
# 몇 개의 샘플 테스트
print("🧪 샘플 테스트 시작...")
test_samples = ds['train'][:5]
for i, sample in enumerate(test_samples):
print(f"\n{'='*80}")
print(f"📝 샘플 {i+1}")
# TriviaQA 샘플 처리
processed_sample = trivia_rag.process_triviaqa_sample(sample, top_k=3)
print(f"❓ 질문: {processed_sample['question']}")
print(f"✅ 답: {processed_sample['answer']}")
print(f"\n📖 Wikipedia 컨텍스트:")
print(f"{processed_sample['wikipedia_context'][:500]}...")
print(f"\n🔍 검색 결과 출처:")
for result in processed_sample['search_results']:
print(f" - {result.metadata['source']} (점수: {result.score:.4f})")
print("\n✅ TriviaQA + Wikipedia RAG 연동 시연 완료!")
except Exception as e:
print(f"❌ 시연 중 오류 발생: {e}")
print("먼저 Wikipedia 데이터셋을 생성해주세요.")
def main_process(config: WikipediaRAGConfig = None):
"""
전체 프로세스를 실행합니다.
Args:
config: 설정 객체
"""
config = config or WikipediaRAGConfig()
print("🚀 Wikipedia RAG + TriviaQA 전체 프로세스 시작!")
print("=" * 80)
try:
# 1단계: Wikipedia 데이터셋 생성
print("\n📚 1단계: Wikipedia RAG 데이터셋 생성")
print("-" * 50)
# 기존 데이터셋 파일 확인
if (os.path.exists(config.pickle_file) and
os.path.exists(config.faiss_index_file)):
print("✅ 기존 데이터셋 파일 발견! 생성 단계를 건너뜁니다.")
user_input = input("새로 생성하시겠습니까? (y/n): ")
if user_input.lower() == 'y':
print("🔄 새로운 데이터셋 생성 중...")
create_wikipedia_rag_dataset(config)
else:
print("📂 기존 데이터셋 사용")
else:
print("🔄 Wikipedia RAG 데이터셋 생성 중...")
create_wikipedia_rag_dataset(config)
# 2단계: 검색 테스트
print("\n🔍 2단계: Wikipedia RAG 검색 테스트")
print("-" * 50)
test_wikipedia_search(config)
# 3단계: TriviaQA 연동 테스트
print("\n🎯 3단계: TriviaQA 연동 테스트")
print("-" * 50)
demonstrate_triviaqa_wikipedia_integration(config)
print("\n🎉 전체 프로세스 완료!")
print("생성된 파일들:")
print(f"- {config.pickle_file}")
print(f"- {config.faiss_index_file}")
print(f"- {config.chromadb_path}/ (ChromaDB 데이터)")
except Exception as e:
print(f"❌ 프로세스 중 오류 발생: {e}")
import traceback
traceback.print_exc()
def quick_test():
"""빠른 테스트를 실행합니다."""
print("⚡ 빠른 테스트 시작!")
print("이 테스트는 5개의 Wikipedia 페이지만 사용합니다.")
print("=" * 60)
# 테스트 데이터셋 생성
dataset = create_wikipedia_rag_dataset(is_test=True)
# 검색 테스트
print("\n🔍 검색 테스트:")
test_wikipedia_search(
pickle_file="test_wikipedia_rag.pkl",
faiss_index_file="test_wikipedia_faiss.index"
)
# TriviaQA 연동 테스트
print("\n🎯 TriviaQA 연동 테스트:")
demonstrate_triviaqa_wikipedia_integration(
pickle_file="test_wikipedia_rag.pkl",
faiss_index_file="test_wikipedia_faiss.index"
)
print("\n✅ 빠른 테스트 완료!")
return dataset
# =====================================================================
# 메인 실행부
# =====================================================================
if __name__ == "__main__":
# NLTK 데이터 설정
setup_nltk()
print("🎯 Wikipedia RAG System 실행 준비 완료!")
print("=" * 60)
print("사용 가능한 함수들:")
print("1. main_process() - 전체 프로세스 실행")
print("2. quick_test() - 빠른 테스트 실행")
print("3. create_wikipedia_rag_dataset() - 데이터셋 생성")
print("4. test_wikipedia_search() - 검색 테스트")
print("5. demonstrate_triviaqa_wikipedia_integration() - TriviaQA 연동 테스트")
print("=" * 60)
# 사용자 입력 받기
print("\n어떤 작업을 실행하시겠습니까?")
print("1: 전체 프로세스")
print("2: 빠른 테스트")
print("3: 종료")
choice = input("선택 (1-3): ").strip()
if choice == '1':
main_process()
elif choice == '2':
quick_test()
elif choice == '3':
print("👋 종료합니다.")
else:
print("❌ 잘못된 선택입니다.")Wikipedia RAG System - 사용 가이드
📋 개요
이 시스템은 Wikipedia 데이터를 활용한 RAG(Retrieval-Augmented Generation) 시스템으로, TriviaQA 질문-답변 시스템을 위한 컨텍스트를 제공합니다.
🏗️ 시스템 구조
1. 핵심 클래스들
WikipediaRAGDataset
- 목적: Wikipedia 데이터 수집, 청킹, 임베딩 생성
- 주요 메서드:
- get_wikipedia_content(): Wikipedia 페이지 수집
- chunk_documents(): 문서 청킹
- create_embeddings(): 임베딩 생성
- create_full_dataset(): 전체 파이프라인 실행
WikipediaRAGRetriever
- 목적: 저장된 데이터를 로드하여 검색 기능 제공
- 주요 메서드:
- search_similar_documents(): FAISS 기반 검색
- search_with_chromadb(): ChromaDB 기반 검색
- pretty_print_results(): 검색 결과 출력
TriviaQAWithWikipediaRAG
- 목적: TriviaQA 데이터셋과 Wikipedia RAG 연동
- 주요 메서드:
- get_context_for_question(): 질문에 대한 컨텍스트 검색
- process_triviaqa_sample(): TriviaQA 샘플 처리
2. 설정 클래스 (WikipediaRAGConfig)
시스템의 모든 설정을 중앙에서 관리합니다:
@dataclass
class WikipediaRAGConfig:
# 모델 설정
embedding_model_name: str = "all-MiniLM-L6-v2"
# 청킹 설정
chunk_size: int = 512
chunk_overlap: int = 50
# 파일 경로
pickle_file: str = "wikipedia_rag_dataset.pkl"
faiss_index_file: str = "wikipedia_faiss.index"
# 기타 설정들...
🚀 사용 방법
1. 빠른 시작
# 1. 빠른 테스트 (5개 페이지만 사용)
quick_test()
# 2. 전체 프로세스 실행
main_process()
2. 커스텀 설정으로 사용
# 커스텀 설정 생성
config = WikipediaRAGConfig(
embedding_model_name="all-mpnet-base-v2", # 더 좋은 모델 사용
chunk_size=1024, # 더 큰 청크
chunk_overlap=100,
pickle_file="custom_dataset.pkl"
)
# 커스텀 Wikipedia 페이지 리스트
custom_titles = [
"Machine Learning",
"Deep Learning",
"Neural Networks",
"Computer Vision",
"Natural Language Processing"
]
# 데이터셋 생성
dataset = create_wikipedia_rag_dataset(
config=config,
titles=custom_titles
)
3. 단계별 사용
# 1단계: 데이터셋 생성
dataset = WikipediaRAGDataset(config)
contents = dataset.get_wikipedia_content(titles)
chunks = dataset.chunk_documents(contents)
embeddings = dataset.create_embeddings(chunks)
dataset.save_to_pickle()
dataset.create_faiss_index()
# 2단계: 검색기 초기화
retriever = WikipediaRAGRetriever(config)
# 3단계: 검색 수행
results = retriever.search_similar_documents("What is machine learning?")
# 4단계: TriviaQA 연동
trivia_rag = TriviaQAWithWikipediaRAG(retriever)
processed_sample = trivia_rag.process_triviaqa_sample(sample)
🔧 커스터마이징
1. 새로운 임베딩 모델 추가
# 설정에서 모델 변경
config = WikipediaRAGConfig(
embedding_model_name="sentence-transformers/all-MiniLM-L12-v2"
)
# 또는 직접 모델 교체
dataset = WikipediaRAGDataset(config)
dataset.embedding_model = SentenceTransformer("custom-model-name")
2. 새로운 텍스트 분할 방법
from langchain.text_splitter import TokenTextSplitter
# 데이터셋 초기화 후 분할기 교체
dataset = WikipediaRAGDataset(config)
dataset.text_splitter = TokenTextSplitter(
chunk_size=512,
chunk_overlap=50
)
3. 새로운 검색 방법 추가
class CustomWikipediaRAGRetriever(WikipediaRAGRetriever):
def search_with_custom_method(self, query: str, top_k: int = 5):
# 커스텀 검색 로직 구현
# 예: 키워드 기반 검색, 하이브리드 검색 등
pass
4. 새로운 데이터 소스 추가
class MultiSourceRAGDataset(WikipediaRAGDataset):
def get_additional_content(self, sources: List[str]):
# 다른 데이터 소스에서 콘텐츠 가져오기
# 예: 웹 크롤링, PDF 파싱, 데이터베이스 등
pass
📊 성능 최적화
1. 배치 크기 조정
config = WikipediaRAGConfig(
embedding_batch_size=64 # GPU 메모리에 따라 조정
)
2. 청킹 전략 최적화
# 작은 청크 (빠른 검색, 낮은 메모리)
config = WikipediaRAGConfig(
chunk_size=256,
chunk_overlap=25
)
# 큰 청크 (더 많은 컨텍스트, 높은 메모리)
config = WikipediaRAGConfig(
chunk_size=1024,
chunk_overlap=100
)
3. 인덱스 최적화
# FAISS 인덱스 타입 변경
def create_custom_faiss_index(embeddings):
dimension = embeddings.shape[1]
# IVF 인덱스 (더 빠른 검색)
quantizer = faiss.IndexFlatIP(dimension)
index = faiss.IndexIVFFlat(quantizer, dimension, 100)
# 학습 및 추가
index.train(embeddings)
index.add(embeddings)
return index
🐛 문제 해결
1. 메모리 부족
- 청크 크기 줄이기
- 배치 크기 줄이기
- 처리할 페이지 수 줄이기
2. 검색 성능 향상
- 더 강력한 임베딩 모델 사용
- 청킹 전략 개선
- 하이브리드 검색 구현
3. 새로운 도메인 적응
- 도메인 특화 Wikipedia 페이지 추가
- 도메인 특화 임베딩 모델 사용
- 커스텀 전처리 파이프라인 구현
📁 파일 구조
wikipedia_rag_system.py # 메인 시스템 파일
requirements.txt # 필요한 패키지
environment.yml # 아나콘다 환경 설정
README.md # 기본 설명서
USAGE_GUIDE.md # 이 가이드 파일
# 생성되는 파일들:
wikipedia_rag_dataset.pkl # 데이터셋
wikipedia_faiss.index # FAISS 인덱스
chroma_db/ # ChromaDB 데이터
🔬 실험 및 평가
1. 검색 성능 평가
def evaluate_search_performance(retriever, test_queries, ground_truth):
"""검색 성능 평가 함수"""
results = []
for query, truth in zip(test_queries, ground_truth):
search_results = retriever.search_similar_documents(query)
# 정확도, 재현율 등 계산
precision = calculate_precision(search_results, truth)
recall = calculate_recall(search_results, truth)
results.append({
'query': query,
'precision': precision,
'recall': recall
})
return results
2. 다양한 모델 비교
models_to_test = [
"all-MiniLM-L6-v2",
"all-mpnet-base-v2",
"sentence-transformers/all-MiniLM-L12-v2"
]
for model_name in models_to_test:
config = WikipediaRAGConfig(embedding_model_name=model_name)
dataset = create_wikipedia_rag_dataset(config)
# 성능 평가...
🤝 확장 가능성
이 시스템은 다음과 같은 방향으로 확장 가능합니다:
- 다중 언어 지원: 한국어, 일본어 등 다른 언어 Wikipedia 지원
- 실시간 업데이트: Wikipedia 변경사항 실시간 반영
- 질문 생성: 주어진 컨텍스트에서 자동 질문 생성
- 답변 생성: 검색된 컨텍스트를 바탕으로 답변 생성
- 웹 인터페이스: 사용자 친화적인 웹 UI 제공
이 가이드를 참고하여 여러분의 용도에 맞게 시스템을 커스터마이징하시기 바랍니다!
728x90
'인공지능 > 논문 리뷰 or 진행' 카테고리의 다른 글
| MIRAGE: A Metric-Intensive Benchmark for Retrieval-Augmented Generation Evaluation (3) | 2025.07.11 |
|---|---|
| DPR - Dense Passage Retrieval for Open-Domain Question Answering (6) | 2025.07.10 |
| DoGe 관련 논문 조사 6 - Synthetic QA Data, SELF-ROUTE (0) | 2025.07.08 |
| DoGe 관련 논문 조사 5 - LayoutLLM, DocLLM (6) | 2025.07.08 |
| DoGe 관련 논문 조사 4 - LongRefiner, PDFTriage (4) | 2025.07.08 |