Jupyter 에이전트: 노트북을 활용한 LLM 추론 훈련

Hugging Face는 작은 언어 모델을 데이터 과학 에이전트로 훈련시키는 파이프라인을 개발했습니다. 이 에이전트는 Jupyter Notebook 내에서 코드를 실행할 수 있습니다. Kaggle 노트북에서 추출한 정제된 합성 데이터셋과 단순화된 에이전트 스캐폴딩을 결합함으로써, 팀은 Qwen3-4B 모델의 데이터 분석 벤치마크 성능을 크게 향상시켰습니다.

데이터 과학 에이전트를 위한 DABStep 벤치마크

에이전트 기반 데이터 과학 역량의 진전을 측정하기 위해 Hugging Face는 DABStep 벤치마크를 활용합니다. 이 벤치마크는 데이터셋을 제공하고 모델이 비정형적인 데이터 질문에 답하도록 요구함으로써 현실적인 작업을 평가합니다. 현재 LLM에게는 매우 어려운 과제로, 예를 들어 Claude 4 Sonnet은 어려운 작업에서 20% 미만의 정확도를 보입니다.

에이전트 스캐폴딩 최적화

에이전트 성능은 스캐폴딩—모델 행동을 유도하는 코드와 프롬프트 구조—에 크게 좌우됩니다. Hugging Face는 Jupyter Agent의 스캐폴딩을 가볍고 모델 요구에 맞게 재구성했습니다:

  • 단순화된 아키텍처: 스캐폴딩을 외부 의존성 없이 약 200줄의 코드로 축소했습니다.
  • 도구 통합: 에이전트는 code execution과 final_answer 두 가지 주요 도구를 사용해 while 루프 안에서 동작합니다.
  • 성능 향상: 이 최적화만으로도 DABStep 쉬운 파트의 정확도가 44.4%에서 59.7%로 상승했습니다.

Jupyter Agent 데이터셋 파이프라인

소형 모델을 개선하기 위해 Hugging Face는 Datatrove를 이용해 약 2TB의 원시 Kaggle 노트북을 고품질 훈련용 51k개의 합성 노트북(약 0.2B 토큰)으로 변환하는 다단계 파이프라인을 구축했습니다.

1. 중복 제거 및 데이터 수집

원시 Kaggle 노트북의 중복을 제거해 용량을 2TB에서 250GB로 줄였습니다. 이후 kagglehub 패키지를 사용해 약 5TB의 연관 데이터셋을 다운로드했으며, 모델 체크포인트, 멀티모달 코퍼스, 10GB를 초과하는 데이터셋은 필터링해 E2B 샌드박스에 맞게 조정했습니다.

2. 교육적 점수 매기기 및 필터링

Qwen3-32B를 활용해 노트북을 명료성 및 교육적 가치 기준으로 1–5점의 "edu scoring" 시스템으로 평가했습니다. 이를 통해 약 70%의 노트북을 제거했으며, 추가로 LLM 기반 필터링으로 20%를 더 제거해 데이터 분석과 무관하거나 데이터셋을 사용하지 않은 노트북을 제외했습니다.

3. 합성 QA 및 트레이스 생성

검증 가능한 훈련 데이터를 만들기 위해 파이프라인은 실제 노트북 트레이스를 기반으로 합성 질문‑답변(QA) 쌍을 생성했습니다.

  • QA 생성: Qwen3-32B가 자연스러운 질문과 답변을 생성하고, 두 번째 LLM이 이를 검증해 환상을 방지했습니다.
  • 트레이스 생성: Qwen-3-Coder-480B-A35B-Instruct를 사용해 합성 질문에 답하기 위한 깔끔한 단계별 코드 실행 트레이스를 만들었습니다.
  • 데이터 부족 처리: 데이터셋이 없을 경우 모델에게 상태를 유지하는 Python 인터프리터 역할을 수행하도록 프롬프트를 제공해 실행을 시뮬레이션했습니다.
  • 추론 모방: 코드 셀 사이에 설명을 삽입하도록 comment 필드를 코드 실행 도구 호출의 필수 항목으로 만들었습니다.

훈련 및 결과

Fine-tuning은 TRL 라이브러리를 사용해 Qwen3-4B에 대해 수행되었습니다. 전체 파라미터 미세조정이 PEFT(파라미터 효율적 미세조정)보다 우수했으며, assistant_loss_only=True 옵션이 성능을 추가로 끌어올렸습니다.

성능 향상

합성 데이터셋으로 훈련한 결과 DABStep 쉬운 파트에서 모델 성능이 크게 개선되었습니다:

모델 에폭 수 DABStep (Easy)
Qwen-3-4B-Instruct-2507 (Base) 0 38.67%
Qwen-3-4B-Instruct-2507 (Our Scaffolding) 0 52.78%
Qwen-3-4B-Instruct-2507 2 63.89%
Qwen-3-4B-Instruct-2507 3 73.61%
Qwen-3-4B-Instruct-2507 5 75%
Qwen-3-4B-Instruct-2507 7 70.83%

최종 튜닝된 Qwen-4B 모델은 기본 모델 대비 쉬운 점수에서 최대 36% 향상, 스캐폴딩 적용 기본 모델 대비 22% 향상을 달성했으며, 어려운 점수에서도 개선을 보였습니다.

오픈 소스 공개

Hugging Face는 커뮤니티에 다음 리소스를 공개했습니다:

  • Jupyter Agent 데이터셋: 51k 합성 노트북 컬렉션.
  • 튜닝된 모델: jupyter-agent-qwen3-4b-instruct와 jupyter-agent-qwen3-4b-thinking.

향후 방향

팀은 더 어려운 다단계 질문 생성, 정제된 트레이스 양 확대, 지식 증류 탐색, 기존 검증 가능한 QA 설정을 활용한 강화 학습(RL) 환경 구축 등 여러 개선 방안을 제시했습니다.

Sources