lotte-evaluator.py

import argparse
import json
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path
from threading import Lock
from typing import Dict, Union

import pandas as pd
from openai import AzureOpenAI, OpenAI

from templates import JUDGE_TEMPLATE, LOTTE_PROMPT_STRATEGY

# Constants
TIME_START = datetime.now().strftime("%Y%m%d_%H%M%S")
LOCK = Lock()

AZURE_ENDPOINT = os.environ.get("AZURE_ENDPOINT", None)
AZURE_DEPLOYMENT_NAME = os.environ.get("AZURE_DEPLOYMENT_NAME", None)
AZURE_API_VERSION = os.environ.get("AZURE_API_VERSION", None)
USE_AZURE_OPENAI = AZURE_ENDPOINT is not None and AZURE_DEPLOYMENT_NAME is not None and AZURE_API_VERSION is not None


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-o", "--model-output-dir", help="Model Output Directory", required=True)
    parser.add_argument("-k", "--openai-api-key", help="OpenAI API Key", required=True)
    parser.add_argument("-j", "--judge-model", help="Judge Model", default="gpt-4-1106-preview")
    parser.add_argument("-t", "--threads", help="Thread count", default=42, type=int)
    parser.add_argument("-m","--model",help=" : write huggingface model name to evaluate",default="LDCC/Chat-Mistral-Nemo-12B-32k",required=True)
    parser.add_argument("--azure", help="Use Azure OpenAI", action="store_true")
    return parser.parse_args()


def create_openai_client(api_key: str):
    return OpenAI(api_key=api_key)


def create_azure_openai_client(api_key: str):
    return AzureOpenAI(
        azure_endpoint=AZURE_ENDPOINT,
        api_key=api_key,
        api_version=AZURE_API_VERSION,
    )


def create_answers(
    client, model_output, judge_model, i=0
) -> Dict[str, Union[str, float]]:
    model_questions = model_output["questions"]
    model_outputs = model_output["outputs"]
    model_references = model_output["references"]
    model_instruct = model_output["category"]

    prompt = (
        f"아래의 내용을 주어진 평가 기준들을 충실히 반영하여 평가해라. 특히 모델 답변이 언어 요구사항을 준수하는지 반드시 확인해야 한다.\n\n"
        f"**Instruct** \n{LOTTE_PROMPT_STRATEGY[model_instruct]}\n\n"
        f"**Question**\n{model_questions[0]}"
    )
    if model_references and model_references[0]:
        prompt += f"\n\n**Additional Reference**\n{model_references[0]}"

    prompt += f"\n\n**Model's Response**\n{model_outputs[0]}"
    prompt += "\n\n[[대화 종료. 평가 시작.]]"

    try:
        if USE_AZURE_OPENAI:  # azure
            response = client.chat.completions.create(
                model=AZURE_DEPLOYMENT_NAME,
                temperature=0.0,
                n=1,
                messages=[
                    {
                        "role": "system",
                        "content": JUDGE_TEMPLATE["lotte_eval_template"],
                    },
                    {"role": "user", "content": prompt},
                ],
            )
        else:  # openai api
            response = client.chat.completions.create(
                model=judge_model,
                temperature=0.0,
                n=1,
                messages=[
                    {
                        "role": "system",
                        "content": JUDGE_TEMPLATE["lotte_eval_template"],
                    },
                    {"role": "user", "content": prompt},
                ],
            )

        content = response.choices[0].message.content
        judge_message_match = re.search(r"평가:(.*?)점수:", content.replace("*", ""), re.DOTALL)
        judge_message = judge_message_match.group(1).strip() if judge_message_match else "No judge message found"
        judge_score_match = re.search(r"점수:\s*(\d+(\.\d+)?)", content.replace("*", ""))
        if judge_score_match:
            judge_score = float(judge_score_match.group(1))
        else:
            raise ValueError("No score found in response")

        return {"judge_message": judge_message, "judge_score": judge_score}

    except Exception as e:
        print("Error. Retrying after 20 sec", e)
        time.sleep(20)

        if i > 3:
            print("Impossible prompt, aborting..!")
            return {
                "judge_message": "Impossible to judge due to repetition.",
                "judge_score": 0.0,
            }
        i += 1
        return create_answers(client, model_output, judge_model, i)


def process_item(client, row, judge_model, output_file):
    query_single = create_answers(client, row, judge_model)

    row["query_single"] = query_single
    row = row.to_dict()

    with LOCK:
        with output_file.open("a", encoding="utf-8-sig") as f:
            f.write(json.dumps(row, ensure_ascii=False))
            f.write("\n")


def process_file(client, file_path: Path, output_dir: Path, judge_model, threads: int, args):
    print(f"- 현재 Processing : {file_path}")
    df_model_outputs = pd.read_json(file_path, lines=True)

    output_file = output_dir / file_path.relative_to(args.model_output_dir)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    with ThreadPoolExecutor(max_workers=threads) as executor:
        for row in df_model_outputs.iterrows():
            executor.submit(process_item, client, row[1], judge_model, output_file)

def is_hidden(filepath: Path) -> bool:
    return any(part.startswith(".") for part in filepath.parts)


def main():
    args = get_args()
    if args.azure:
        client = create_azure_openai_client(args.openai_api_key)
    else:
        client = create_openai_client(args.openai_api_key)

    input_dir = Path(args.model_output_dir)
    output_dir = Path(f"./evaluated/{args.model}")

    # lotte_single_turn.jsonl 파일만 찾도록 수정
    json_file = input_dir / "lotte_single_turn.jsonl"
    if not json_file.exists():
        print(f"Error: {json_file} not found")
        return

    output_file_path = output_dir / json_file.relative_to(input_dir)
    if output_file_path.exists():
        print(f"추가된 데이터로 롯데 벤치마크를 진행합니다. : {json_file}")
    process_file(client, json_file, output_dir, args.judge_model, args.threads, args)

if __name__ == "__main__":
    main()