Mistral API로 시스템 수준 가드레일 구현하기

Mistral API로 시스템 수준 가드레일 구현하기 (System-Level Guardrails)

Mistral의 분류기 기반 컨텐츠 중재(moderation) 서비스를 활용해 시스템 수준의 가드레일을 구현하는 방법을 다루는 문서예요. 텍스트와 대화 콘텐츠를 9개 카테고리로 분류하고, 생성된 응답을 점수화·정렬해 안전한 응답을 고르는 전체 파이프라인을 배워요.

출처: 문서

본문

Mistral은 Ministral 8B 24.10 기반의 분류기 모델로 구동되는 중재 서비스를 제공해요. 품질이 높고 빠르면서도 뛰어난 성능을 내고, 다음 두 가지를 모두 중재할 수 있습니다:

• 텍스트 콘텐츠 (Text content)
• 대화 콘텐츠 (Conversational content)

안전(Safeguarding)과 중재에 대한 자세한 내용은 [관련 문서]를 참고하시면 돼요.

개요 (Overview)

이 튜토리얼에서는 Mistral 클라이언트를 설정하고, 응답을 생성하고, 대화를 중재하며, 결과를 시각화하는 전체 과정을 안내해요. 텍스트나 대화 데이터를 9개 카테고리로 쉽게 분류할 수 있으며, 대화 데이터의 경우 마지막 사용자 메시지가 분류 대상이 됩니다.

분류 카테고리는 다음과 같아요:

• Sexual
• Hate and Discrimination
• Violence and Threats
• Dangerous and Criminal Content
• Self-harm
• Health
• Financial
• Law
• PII (Personally Identifiable Information)

테스트를 위해 Hugging Face와 GitHub의 데이터셋을 사용할 예정이에요.

Step 1: 설정 (Setup)

먼저 클라이언트를 설정해 볼게요. 이 쿡북은 v1.2.3에서 테스트됐어요.

!pip install mistralai

API 키를 추가하고 클라이언트를 만들어요 (API 키는 [여기]에서 생성할 수 있어요).

from mistralai.client import Mistral

api_key = "API_KEY"

client = Mistral(api_key=api_key)

Step 2: 응답 생성 (Generate Responses)

어떤 Mistral 모델로부터든 응답을 생성하는 함수를 만들어요.

def generate_responses(client: Mistral, user_prompt: str, num_generations: int) -> list:
    """
    Generate responses from the Mistral model.

    Args:
        client (Mistral): The Mistral client instance.
        user_prompt (str): The user prompt.
        num_generations (int): The number of generations to produce.

    Returns:
        list: A list of generated responses.
    """
    chat_response = client.chat.complete(
        n=num_generations,
        model="mistral-large-latest",
        temperature=0.3, # Adds randomness to generate diverse outputs
        messages=[{"role": "user", "content": user_prompt}],
    )
    responses = chat_response.choices
    assert len(responses) == num_generations
    return responses

# Quick test
test_prompt = "Tell me a short story."
test_responses = generate_responses(client, test_prompt, 5)
test_str = '\n- '.join([response.message.content for response in test_responses])
print(f"Generated Responses:\n- {test_str}")

이 함수는 사용자 프롬프트와 생성 개수를 입력받아, 어떤 Mistral 모델로부터든 생성된 응답 리스트를 돌려줘요. 여기서는 mistral-large-latest를 선택했어요.

보통 각 응답은 조금씩 달라지는데, temperature 같은 샘플링 설정에 따라 서로 다를 수도, 덜 다를 수도 있어요. 응답 생성에는 client.chat.complete 메서드가 사용됩니다.

Step 3: 대화 중재 (Moderate Conversation)

Mistral 중재 API를 사용해 대화를 중재하는 함수를 만들어요.

def moderate_conversation(client: Mistral, user_prompt: str, response: str) -> dict:
    """
    Moderate the conversation using the Mistral moderation API.

    Args:
        client (Mistral): The Mistral client instance.
        user_prompt (str): The user prompt.
        response (str): The assistant response.

    Returns:
        dict: The moderation results.
    """
    response = client.classifiers.moderate_chat(
        model="mistral-moderation-latest",
        inputs=[
            {"role": "user", "content": user_prompt},
            {"role": "assistant", "content": response},
        ],
    )
    return response

# Quick test
test_moderation = moderate_conversation(client, test_prompt, test_responses[0].message.content)
from pprint import pprint
pprint(dict(test_moderation))

이 함수는 사용자 프롬프트와 어시스턴트 응답을 입력받아 중재 결과를 돌려줘요.

Step 4: 응답 점수화 및 정렬 (Score and Sort Responses)

중재 결과를 바탕으로 응답에 점수를 매기고 정렬하는 함수를 만들어요.

def score_and_sort_responses(client: Mistral, user_prompt: str, responses: list, threshold: float = 0.2) -> tuple:
    """
    Score and sort the responses based on the moderation results.

    Args:
        client (Mistral): The Mistral client instance.
        user_prompt (str): The user prompt.
        responses (list): A list of generated responses.
        threshold (float): if max(moderation_score) is above this threshold
        we will return a preformulated response to the user. This threshold
        should be customized by the user depending on their use case and
        safety standards.

    Returns:
        tuple: The final response and the list of scores.
    """
    scores_list = []
    for response in responses:
        scores = moderate_conversation(client, user_prompt, response.message.content)
        scores = scores.results[0]
        category_scores = scores.category_scores
        # The user should customize which policies they include here.
        max_score = max(category_scores.values())
        scores_list.append((response, max_score))

    # Sort the list of responses based on the maximum scores in ascending order,
    # making use of all 9 categories (to be adjusted by users).
    sorted_responses = sorted(scores_list, key=lambda x: x[1])

    lowest_score = sorted_responses[0][1]
    if lowest_score >= threshold:
        final_response = "I'm sorry I cannot assist with this request."
    else:
        final_response = sorted_responses[0][0]

    return final_response, sorted_responses

# Quick test
final_response, sorted_responses = score_and_sort_responses(client, test_prompt, test_responses)
print(f"Final Response: {final_response.message.content if isinstance(final_response, dict) else final_response}")
print(f"Sorted Responses: {[(response.message.content, score) for response, score in sorted_responses]}")

이 함수는 사용자 프롬프트와 생성된 응답 리스트를 입력받아 최종 응답과 점수 리스트를 돌려줘요. 중재 결과로 각 응답에 점수를 매기고, 최고 점수 기준 오름차순으로 정렬합니다. 만약 가장 낮은 점수가 특정 임계값(threshold)보다 높으면 기본 안전 응답을 반환해요.

Step 5: 결과 시각화 (Visualize Results)

중재 결과를 시각화하는 함수를 만들어요.

def visualize_responses(user_prompt: str, responses: list, scores_list: list, final_response: str) -> None:
    """
    Visualize the responses with their scores and indicate whether they were chosen or not.

    Args:
        user_prompt (str): The user prompt.
        responses (list): A list of generated responses.
        scores_list (list): A list of tuples containing responses and their scores.
        final_response (str): The final chosen response.
    """
    print("=" * 50)
    print(f"User Prompt: {user_prompt}\n")
    print("### Responses and Scores:\n")

    for response, score in scores_list:
        chosen = "Chosen" if response.message.content == final_response.message.content else "Not Chosen"
        print(f"Response: {response.message.content}")
        print(f"Highest Unsafe Score: {score:.2f}")
        print(f"Status: {chosen}\n")
        print("-" * 50)

# Quick test
visualize_responses(test_prompt, test_responses, sorted_responses, final_response)

이 함수는 사용자 프롬프트, 응답 리스트, 점수 리스트, 최종 응답을 입력받아 각 응답과 점수를 출력하고 선택됐는지 여부를 표시해요. 선택되지 않은 경우 대신 기본 안전 응답이 선택됐음을 알 수 있어요.

Step 6: 데이터셋 함수 (Dataset Function)

전체 과정을 데이터셋에 대해 실행하는 함수를 만들어요.

def run(input_dataset: list) -> None:
    for user_prompt in input_dataset:
        responses = generate_responses(client, user_prompt, 3) # Here we arbitrary decided to generate 3 variations of responses
        final_response, scores_list = score_and_sort_responses(client, user_prompt, responses)
        visualize_responses(user_prompt, responses, scores_list, final_response)

이 함수는 입력 데이터셋을 받아 데이터셋의 각 사용자 프롬프트에 대해 전체 과정을 실행해요. 응답을 생성하고, 점수화·정렬하며, 결과를 시각화합니다.

Step 7: 데이터셋 로드 (Load Datasets)

테스트용 데이터셋을 Hugging Face와 GitHub에서 로드해요.

!pip install datasets

import pandas as pd
from datasets import load_dataset
import random

# Load toxic chat dataset from Hugging Face, having both safe and unsafe examples
toxic_chat_dataset = load_dataset('lmsys/toxic-chat', 'toxicchat0124')

# Load harmful strings dataset from GitHub, with mostly unsafe examples
harmful_strings_url = "https://raw.githubusercontent.com/llm-attacks/llm-attacks/main/data/advbench/harmful_strings.csv"
harmful_strings_df = pd.read_csv(harmful_strings_url)

# Combine datasets
combined_dataset = toxic_chat_dataset['train']['user_input'] + harmful_strings_df['target'].tolist()

# Suffle them
seed = 42
random.seed(seed)
random.shuffle(combined_dataset)

• lmsys/toxic-chat 데이터셋은 안전·불안전 예시를 모두 포함하고 있고, harmful_strings.csv는 대부분 불안전한 예시로 구성돼 있어요. 두 데이터셋을 합친 뒤 셔플합니다.

Step 8: 실행 (Run)

5개 샘플로 실행하고 결과를 시각화해요.

run(combined_dataset[:5])

이 코드는 결합된 데이터셋의 처음 5개 샘플에 대해 함수를 실행하고 결과를 시각화해요. 보시다시피, 모든 응답이 임계값과 생성 개수에 따라 중재된 것을 확인할 수 있어요.

더 알아보기 (Learn more)

• Mistral 중재(Moderation) API 문서
• Mistral 쿡북 목록 (Cookbooks)
• lmsys/toxic-chat 데이터셋