微调谷歌开源Gemini Flash模型实现PII脱敏实战 原创

A student’s assessment was found on device bearing IMEI: 06–184755–866851–3. The document falls under the various topics discussed in our Optimization curriculum. Can you please collect it?

A student’s assessment was found on device bearing IMEI: [PHONEIMEI]. The document falls under the various topics discussed in our [JOBAREA] curriculum. Can you please collect it?

pip install transformers datasets evaluate google-generativeai python-dotenv sacrebleu

#导入必需的库
from datasets import load_dataset
from google.generativeai.types import HarmCategory, HarmBlockThreshold
# Define GOOGLE_API_KEY as a global variable
# 加载和分割数据集函数
def load_data(train_size: int, test_size: int):
    """
    加载pii-masking-200k数据集并把它拆分成训练与测试两种类型。
    参数：
        train_size: 训练数据集的大小
        test_size: 测试数据集的大小
    返回：
       一个包含训练数据集和测试数据集的元组。
    """
    dataset = load_dataset("ai4privacy/pii-masking-200k")
    dataset = dataset["train"].train_test_split(test_size=test_size, seed=42)
    train_d = dataset["train"].select(range(train_size))
    test_d = dataset["test"]
    return train_d, test_d

import google.generativeai as genai
from google.generativeai.types.content_types import ContentDict
from google.generativeai.types import HarmCategory, HarmBlockThreshold

from concurrent.futures import ThreadPoolExecutor
import evaluate

safety_settings = {
    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
}
SYS_PROMPT = (
    "Substitute all PII in this text for a generic label like [FIRSTNAME] (Between square brackets)\n"
    "Labels to substitute are PREFIX, FIRSTNAME, LASTNAME, DATE, TIME, "
    "PHONEIMEI, USERNAME, GENDER, CITY, STATE, URL, JOBAREA, EMAIL, JOBTYPE, "
    "COMPANYNAME, JOBTITLE, STREET, SECONDARYADDRESS, COUNTY, AGE, USERAGENT, "
    "ACCOUNTNAME, ACCOUNTNUMBER, CURRENCYSYMBOL, AMOUNT, CREDITCARDISSUER, "
    "CREDITCARDNUMBER, CREDITCARDCVV, PHONENUMBER, SEX, IP, ETHEREUMADDRESS, "
    "BITCOINADDRESS, MIDDLENAME, IBAN, VEHICLEVRM, DOB, PIN, CURRENCY, "
    "PASSWORD, CURRENCYNAME, LITECOINADDRESS, CURRENCYCODE, BUILDINGNUMBER, "
    "ORDINALDIRECTION, MASKEDNUMBER, ZIPCODE, BIC, IPV4, IPV6, MAC, "
    "NEARBYGPSCOORDINATE, VEHICLEVIN, EYECOLOR, HEIGHT, SSN, language"
)
#计算零样本设置的函数
def evaluate_zero_shot(train_data, test_data, model_name="gemini-1.5-flash"):
    """
    评估该模型的零样本性能。
    参数：
        train_data: 训练数据集（不用于零样本情况）。
        test_data: 测试数据集。
        model_name: 要使用的模型的名称。
    返回值 ：
        零样本设置的SacreBLEU分数
    """
    model = genai.GenerativeModel(model_name)
    def map_zero_shot(text):
        messages = [
            ContentDict(
                role="user",
                parts=[f"{SYS_PROMPT}\nText: {text}"],
            ),
        ]
        response = model.generate_content(messages, safety_settings=safety_settings)
        try:
            return response.text
        except ValueError:
            print(response)
            return ""
    with ThreadPoolExecutor(max_workers=4) as executor:
        predictions = list(
            executor.map(
                map_zero_shot,
                [example["source_text"] for example in test_data],
            )
        )
    references = [[example["target_text"]] for example in test_data]
    sacrebleu = evaluate.load("sacrebleu")
    sacrebleu_results = sacrebleu.compute(
        predictions=predictions, references=references
    )
    print(f"Zero-shot SacreBLEU score: {sacrebleu_results['score']}")
    return sacrebleu_results["score"]

# 评估小样本设置的函数
def evaluate_few_shot(train_data, test_data, model_name="gemini-1.5-flash"):
    """
    评估模型的小样本性能
    参数：
        train_data: 训练数据集
        test_data: 测试数据集
        model_name: 要使用的模型的名称。
    返回值：
        小样本设置的SacreBLEU分数
    """
    model = genai.GenerativeModel(model_name)
    def map_few_shot(text, examples):
        messages = [
            ContentDict(
                role="user",
                parts=[SYS_PROMPT],
            )
        ]
        for example in examples:
            messages.append(
                ContentDict(role="user", parts=[f"Text: {example['source_text']}"]),
            )
            messages.append(
                ContentDict(role="model", parts=[f"{example['target_text']}"])
            )
        messages.append(ContentDict(role="user", parts=[f"Text: {text}"]))
        response = model.generate_content(messages, safety_settings=safety_settings)
        try:
            return response.text
        except ValueError:
            print(response)
            return ""
    few_shot_examples = train_data.select(range(3))
    with ThreadPoolExecutor(max_workers=4) as executor:
        predictions = list(
            executor.map(
                lambda example: map_few_shot(example["source_text"], few_shot_examples),
                test_data,
            )
        )
    references = [[example["target_text"]] for example in test_data]
    sacrebleu = evaluate.load("sacrebleu")
    sacrebleu_results = sacrebleu.compute(
        predictions=predictions, references=references
    )
    print(f"3-shot SacreBLEU score: {sacrebleu_results['score']}")
    return sacrebleu_results["score"]

def finetune(train_data, finetune_size, model_name="gemini-1.5-flash"):
    """
    调优模型

    参数:
        train_data: 训练数据集。
        finetune_size:用于微调的样本数量。
        model_name: 要用于进行微调的基本模型的名称。
    返回值:
        已调优的模型的名称。
    """
    base_model = f"models/{model_name}-001-tuning"
    tuning_data = [
        {
            "text_input": f"{SYS_PROMPT}\nText: {example['source_text']}",
            "output": example["target_text"],
        }
        for example in train_data.select(range(finetune_size))
    ]
    print(len(tuning_data))
    operation = genai.create_tuned_model(
        display_name=f"tuned-{finetune_size}",
        source_model=base_model,
        epoch_count=2,
        batch_size=4,
        learning_rate=0.0001,
        training_data=tuning_data,
    )

import google.generativeai as genai

for model_info in genai.list_tuned_models():
    print(model_info.name)
    print(model_info)