데이터브릭스에서 손쉽게 Text2SQL 성능 개선하기

Text-to-SQL 작업에서 널리 사용되는 벤치마크인 Spider의 상위 10위권으로 LLM을 끌어올리고 싶으신가요? Spider는 LLM이 텍스트 쿼리를 SQL 코드로 얼마나 잘 변환할 수 있는지 평가합니다.

Text-to-SQL에 익숙하지 않은 분들을 위해 설명하자면, 이 변환은 기업이 데이터와 상호 작용하는 방식을 혁신하는 데 중요한 역할을 합니다. SQL 전문가에게 쿼리 작성을 의존하지 않고 일상 언어로 데이터에 대해 질문하고 정확한 답변을 받을 수 있습니다. 이를 통해 데이터에 대한 액세스가 민주화되어 비즈니스 인텔리전스가 향상되고 정보에 입각한 의사결정이 가능해집니다.

Spider 벤치마크는 Text-to-SQL 시스템의 성능을 평가하는 데 널리 인정받는 표준입니다. 이 테스트는 LLM이 자연어 쿼리를 정확한 SQL 문으로 번역하는지 검증하며, 데이터베이스 스키마에 대한 깊은 이해와 구문 및 의미론적으로 정확한 SQL 코드를 생성할 수 있는 능력을 요구합니다.

이 포스트에서는 오픈 소스 Llama3 8B Instruct 모델을 사용하여 하루도 채 안 되는 작업 시간 동안 Spider 개발 데이터 세트에서 79.9%, 테스트 데이터 세트에서 78.9%의 점수를 달성하여 기준점 대비 19점이라는 놀라운 개선을 이룬 방법을 자세히 살펴보겠습니다. 이러한 성능은 현재 정지된 Spider 순위표에서 상위 10위 안에 들어 있는데, 이는 데이터브릭스에서 수행한 전략적인 프롬프트와 미세 조정 덕분입니다.

기준 성능을 위한 제로 샷 프롬프트(Zero-shot Prompting)

먼저 테이블을 생성한 CREATE TABLE 문과 해당 테이블을 사용하여 답변할 질문으로 구성된 매우 간단한 프롬프트 형식을 사용하여, Spider 개발 데이터세트에 대한 Meta Llama 3 8B Instruct의 성능을 평가해 보겠습니다:

이러한 유형의 프롬프트는 프롬프트에 다른 예가 없기 때문에 종종 "제로 샷(zero-shot)"이라고도 합니다. Spider 개발 데이터 세트의 첫 번째 질문에 대해 이 프롬프트 포맷은 다음과 같이 생성됩니다:

이 포맷을 사용하여 개발 데이터 셋에서 Spider 벤치마크를 실행하면 execution accuracy와 greedy decording을 사용하여 측정한 결과 전체 점수가 60.9점이 나옵니다. 즉, 60.9%의 경우 모델이 실행 시 올바른 솔루션을 나타내는 "골드" 쿼리와 동일한 결과를 내는 SQL을 생성한다는 의미입니다.

기준 점수가 설정되었으므로, 미세 조정을 시작하기 전에 Spider 개발 벤치마크 데이터셋에서 기본 모델의 점수를 높이기 위해 다양한 프롬프트 전략을 시도해 보겠습니다.

샘플 행을 사용한 프롬프트

우리가 사용한 첫 번째 프롬프트의 단점 중 하나는 데이터 유형 외에 열의 데이터에 대한 정보가 포함되어 있지 않다는 것입니다. Spider로 모델의 Text-to-SQL 기능을 평가하는 한 논문에 따르면, 프롬프트에 샘플 행을 추가하면 점수가 더 높아지므로 이를 시도해 보겠습니다.

테이블 생성 쿼리에 각 테이블의 처음 몇 개의 행도 포함되도록 위의 프롬프트 형식을 업데이트할 수 있습니다. 앞의 동일한 질문에 대해 업데이트된 프롬프트는 다음과 같습니다:

각 테이블에 샘플 행을 포함하면 전체 점수가 67.0으로 약 6% 포인트 상승합니다:

퓨 샷 (Few-shot) 프롬프트

퓨 샷 프롬프트는 LLM에 사용되는 잘 알려진 전략으로, 수행할 작업을 보여주는 몇 가지 예제를 포함함으로써 올바른 SQL 생성과 같은 작업의 성능을 향상시키는 방법입니다. 제로 샷 프롬프트에서는 스키마를 제공한 다음 바로 질문을 했지만, 퓨 샷 프롬프트에서는 몇 가지 스키마, 질문, 해당 질문에 답하는 SQL을 제공한 다음 이 순서를 몇 번 반복한 후 실제 질문으로 넘어갑니다. 이렇게 하면 일반적으로 제로 샷 프롬프트보다 더 나은 성능을 얻을 수 있습니다.

SQL 생성 작업을 보여주는 좋은 예제는 사실 Spider 학습 데이터셋 자체입니다. 이 데이터셋에서 해당 테이블과 함께 몇 가지 질문의 무작위 샘플을 가져와서 각 질문에 답할 수 있는 SQL을 보여주는 퓨 샷 프롬프트를 만들 수 있습니다. 이전 프롬프트에서 샘플 행을 사용하고 있으므로 이러한 예제에도 샘플 행을 포함하는 사용법을 보여주어야 합니다.

이전의 제로 샷 프롬프트에서 개선할 수 있는 또 다른 사항은 시작 부분에 "시스템 프롬프트(system prompt)"를 추가하는 것입니다. 시스템 프롬프트는 일반적으로 모델에게 수행해야 할 작업에 대한 자세한 안내를 제공하는 데 사용됩니다. 사용자는 모델과 채팅하는 동안 여러 가지 질문을 할 수 있지만, 시스템 프롬프트는 사용자가 질문을 하기 전에 한 번만 제공되므로 기본적으로 채팅 중에 "시스템"이 어떻게 작동해야 하는지에 대한 기대치를 설정할 수 있습니다.

이러한 전략을 염두에 두고, 아래와 같이 상단에 큰 SQL 코멘트 블록으로 표시된 시스템 메시지에 이어 세 가지 예제가 제공되는 퓨 샷 프롬프트를 구성할 수 있습니다:

이 새로운 프롬프트의 결과 점수는 70.8점으로, 이전 점수보다 3.8% 포인트 향상되었습니다. 단순한 프롬프트 전략만으로 시작점 대비 10% 포인트 가까이 점수를 높인 것입니다.

이제 프롬프트 조정으로 인한 수익이 점차 감소하는 시점에 도달한 것 같습니다. 모델을 미세 조정(fine-tune)하여 어떤 추가 이득을 얻을 수 있는지 살펴봅시다.

LoRA로 미세 조정하기

모델을 미세 조정하는 경우 첫 번째 질문은 어떤 학습 데이터를 사용할 것인가입니다. Spider에는 학습 데이터셋이 포함되어 있으므로 여기서부터 시작하면 좋을 것 같습니다. 모델을 미세 조정하기 위해 Standard_NC24ads_A100_v4와 같은 단일 A100 80GB 데이터브릭스 GPU 클러스터에서 모델을 효율적으로 훈련할 수 있도록 QLoRA를 사용할 것입니다. Spider 학습 데이터셋에 있는 7천 개의 레코드를 사용하면 이 작업은 약 4시간 안에 완료할 수 있습니다. 이전 블로그 게시물에서 LoRA를 사용한 미세 조정에 대해 설명한 적이 있습니다. 관심 있는 독자는 해당 포스팅에서 자세한 내용을 확인할 수 있습니다. trl, peft, bitsandbytes 라이브러리를 사용하여 표준 훈련 레시피를 따를 수 있습니다.

Spider에서 학습 레코드를 가져오더라도 모델이 학습할 수 있는 방식으로 학습 데이터를 포매팅해야 합니다. 목표는 스키마(샘플 행 포함), 질문, SQL로 구성된 각 레코드를 단일 텍스트 문자열로 매핑하는 것입니다. 먼저 원시 Spider 데이터셋에 대해 몇 가지 처리를 수행합니다. 원시 데이터로부터 각 레코드가 세 개의 필드, 즉 schema_with_rows, question, query로 구성된 데이터셋을 생성합니다. schema_with_rows 필드는 질문에 해당하는 테이블에서 파생되며, 앞서 퓨 샷 프롬프트에서 사용된 CREATE TABLE 문과 샘플 행의 서식을 따릅니다.

다음으로 토크나이저(tokenizer)를 로드합니다:

처리된 Spider 학습 데이터셋의 각 레코드를 텍스트 문자열로 변환하는 매핑 함수를 정의하겠습니다. 토크나이저의 apply_chat_template를 사용하여 텍스트를 Instruct 모델에서 예상하는 채팅 형식으로 편리하게 포맷할 수 있습니다. 이 형식은 앞서 퓨 샷 프롬프트에 사용한 형식과 정확히 동일하지는 않지만, 이 모델은 프롬프트의 보일러플레이트 포매팅이 약간 다르더라도 충분히 잘 작동할 만큼 일반화되어 있습니다.  

SYSTEM_PROMPT에는 앞서 퓨 샷 프롬프트에서 사용한 것과 동일한 시스템 프롬프트를 사용합니다. USER_MESSAGE_FORMAT의 경우에도 마찬가지로 사용합니다:

이 함수가 정의되면, 처리된 Spider 데이터셋을 이 함수로 변환하여 JSONL 파일로 저장하면 됩니다.

이제 훈련할 준비가 되었습니다. 몇 시간 후 미세 조정된 Llama3 8B Instruct를 받았습니다. 이 새로운 모델에서 우리의 퓨 샷 프롬프트를 다시 실행한 결과 79.9점으로 이전 점수보다 9% 포인트 더 향상되었습니다. 이제 총 점수는 단순 제로 샷 기준점보다 약 19% 포인트 상승했습니다.

이제 Llama3 8B Instruct 모델과 미세 조정된 버전이 Llama3 70B Instruct와 같은 더 큰 모델과 비교하면 어떤지 궁금하실 것입니다. 개발 데이터셋에서 기성 70B 모델과 8개의 A100 40GB GPU를 사용하여 평가 프로세스를 반복하고 그 결과를 아래에 기록했습니다.

예상대로, 기성 모델들(off-the-shelf models)을 비교하면 동일한 프롬프트 형식으로 측정했을 때 70B 모델이 8B 모델을 앞섰습니다. 하지만 놀라운 점은 미세 조정된 Llama3 8B Instruct 모델이 Llama3 70B Instruct 모델보다 3% 포인트 더 높은 점수를 받았다는 점입니다. Text-to-SQL 변환과 같은 특정 작업에 집중할 경우, 미세 조정을 통해 크기가 훨씬 큰 모델과 필적할만한 성능의 작은 모델을 만들 수 있습니다.

모델 서빙 엔드포인트에 배포하기

Llama3는 Mosaic AI Model Serving에서 지원되므로 미세 조정된 Llama3 모델을 엔드포인트에 배포하여 애플리케이션을 구동하는 데 사용할 수도 있습니다. 미세 조정된 모델을 Unity Catalog에 로깅한 다음 UI를 사용하여 엔드포인트를 생성하기만 하면 됩니다. 엔드포인트가 배포되면 공통 라이브러리를 사용하여 쿼리할 수 있습니다.

마무리

우리는 Llama3 8B Instruct 모델에서 Spider 개발 데이터셋에 대한 제로 샷 프롬프트로 여정을 시작하여 60.9점이라는 준수한 점수를 얻었습니다. 이어서 시스템 메시지, 여러 예제 및 샘플 행으로 완성된 퓨 샷 프롬프트로 이를 개선하여 점수를 70.8점까지 끌어 올렸습니다. Spider 학습 데이터셋에서 모델 미세 조정을 통해 Spider dev에서 79.9점, Spider test에서 78.9점으로 인상적인 점수를 얻었습니다. 이는 시작점 대비 19점이나 크게 상승한 수치이며, Llama3 70B Instruct 기본 모델보다 3점이나 앞선 수치로서, 우리 모델의 성능을 보여주었을 뿐만 아니라 Spider 결과에서도 상위 10위 이내의 위치를 차지할 수 있게 했습니다.

Data+AI Summit에 등록하여 오픈 소스 LLM과 데이터 인텔리전스 플랫폼의 강력한 기능을 활용하는 방법에 대해 자세히 알아보세요.

(번역: Youngkyong Ko) Original Post

부록

평가 설정

Generation은 vLLM, greedy decording(temperature 0), A100 80GB GPU 2개, 최대 1024개의 새 토큰을 사용하여 수행했습니다. Generation을 평가하기 위해 Github의 taoyds/test-suite-sql-eval 리포지토리에 있는 test suite를 사용했습니다.

훈련 설정

다음은 미세 조정(fine-tuning) 설정에 대한 구체적인 세부 사항입니다:

제로 샷 프롬프트 예제

다음은 평가에 사용한 개발 데이터셋의 첫 번째 레코드로, 테이블 스키마를 포함하는 제로 샷 프롬프트 형식으로 되어 있습니다. 질문에서 다루는 테이블은 해당 테이블을 생성한 CREATE TABLE 문을 사용하여 표시됩니다.

샘플 행이 포함된 제로 샷 프롬프트 예제

다음은 평가에 사용한 개발 데이터셋의 첫 번째 레코드로, 테이블 스키마와 샘플 행을 포함하는 제로 샷 프롬프트 형식으로 되어 있습니다. 질문에서 다루는 테이블은 해당 테이블을 생성한 CREATE TABLE 문을 사용하여 표현되었습니다. 행은 각 테이블에서 "SELECT * {table_name} LIMIT 3"을 사용하여 선택되었으며 열 이름은 헤더로 표시됩니다.

샘플 행이 포함된 퓨 샷 프롬프트 예제

다음은 평가에 사용한 개발 데이터셋의 첫 번째 레코드로, 테이블 스키마와 샘플 행을 포함하는 퓨 샷 프롬프트 형식으로 형식이 지정되어 있습니다. 질문에서 다루는 테이블은 해당 테이블을 생성한 CREATE TABLE 문을 사용하여 표현되었습니다. 행은 각 테이블에서 "SELECT * {table_name} LIMIT 3"을 사용하여 선택되었으며 열 이름은 헤더로 표시됩니다.