首页

手游

资讯

专题

首页 > 教程攻略 > ai资讯 >Llama3 微调，增强知识图谱关系抽取（附微调数据集和代码）

Llama3 微调，增强知识图谱关系抽取（附微调数据集和代码）

来源：互联网时间：2026-08-03 14:18:15

通过这篇文章，你将掌握两件事：第一，搞清楚如何用微调Llama3的方式，来强化知识图谱中的关系抽取能力；第二，拿到可复现的微调数据集和配套代码。

关系抽取（Relation Extraction，简称RE），说白了就是从非结构化文本里把关系给挖出来——识别出不同实体之间到底有什么样的联系。它和命名实体识别（NER）通常是搭伴干活，是自然语言处理流程中非常基础的一步。

不过，随着大语言模型（LLMs）的兴起，传统的监督方法——先标注实体范围，再分类它们之间的关系——现在正在被基于LLM的方法所增强或彻底取代。

Llama3是生成式AI领域最近一次重磅发布，目前提供了两个版本：8B和70B。400B的版本据说也快了。这些模型在HuggingFace上都能找到。70B这个版本支撑着Meta的新聊天网站Meta.ai，表现已经不输ChatGPT。而8B版本在同体量模型中，也是顶级的水平。

Llama3的架构和Llama2很像，性能的提升主要靠的是数据升级。它配备了更强大的分词器，上下文窗口也扩展了。虽然官方只放出了一小部分数据，但整体上，这确实是个出色的开源模型。

Teacher and Student

70B版本确实能生成让人眼前一亮的输出，但问题在于它的规模太大，在本地跑起来成本高、效率低，不太现实。那怎么办？思路很直接：让老师（Llama3-70B）教学生（Llama3-8B）——教会小模型怎么从非结构化文本里准确抽取关系。

具体操作是这样的：借助Llama3-70B，我们构建了一个专门用于关系抽取的监督微调数据集。然后，用这个数据集去微调Llama3-8B，从而大幅提升它在这项任务上的能力。

如果想在Google Colab里复现整个流程，需要提前准备好两样东西：

HuggingFace token（用来保存微调后的模型，可选）以及Llama3的访问权限。按模型卡片上的说明操作即可。
一个免费的GroqCloud账户（用Google账号就能登录），外加相应的API key。

搞AI模型的，这些账户应该都是标配。如果没有，确实需要认真考虑一下了。

工作区设置

这个项目使用了配备A100 GPU和高内存设置的Google Colab Pro。先说下环境准备。

先把所有必需的库安装好：

!pip install -q groq  
!pip install -U accelerate bitsandbytes datasets evaluate   
!pip install -U peft transformers trl

整个配置过程非常顺利，没有任何依赖问题，也不需要从源码安装`transformers`，即使模型是全新的。

接着，需要给Google Colab授权访问驱动器，并设置工作目录：

# For Google Colab settings  
from google.colab import userdata, drive  

# 这将提示进行授权  
drive.mount('/content/drive')  

# 设置工作目录  
%cd '/content/drive/MyDrive/postedBlogs/llama3RE'

如果想把模型上传到HuggingFace Hub，还需要配置Hub凭据。这一步是可选的，凭证可以存在Google Colab的key里（通过左侧的键按钮访问）。

Hugging Face Hub 设置

from huggingface_hub import login

从Colab secrets中读取HuggingFace token（需具有写入权限）：

HF = userdata.get('HF')

然后登录并上传模型：

login(token=HF,add_to_git_credential=True)

为了方便文件访问，还添加了一些路径变量：

# 为数据文件夹创建一个路径变量
data_path = '/content/drive/MyDrive/postedBlogs/llama3RE/datas/'

# 完整的微调数据集
sft_dataset_file = f'{data_path}sft_train_data.json'

# Data collected from the mini-test
mini_data_path = f'{data_path}mini_data.json'            

# 测试数据包含所有三个输出
all_tests_data = f'{data_path}all_tests.json'
# 调整后的训练数据集
train_data_path = f'{data_path}sft_train_data.json'

# 创建一个路径变量，用于将SFT模型保存到本地
sft_model_path = '/content/drive/MyDrive/llama3RE/Llama3_RE/'

工作区搞定之后，就可以进入正题了：为关系抽取任务构建一个合成数据集。

使用 Llama3-70B 创建合成数据集

目前，有一些公开的关系抽取数据集可用，比如CoNLL04、web_nlg（HuggingFace上可找到），以及AllenAI开发的SciREX。但大多数数据集都带有比较严格的许可协议。

受`web_nlg`数据格式的启发，我们决定自己构建数据集。如果打算后续对模型做微调，这种自建数据集的方式会特别有用。第一步，是准备一批适用于关系抽取任务的短句子。语料库的构建方式可以很灵活。

收集句子

我们将使用databricks-dolly-15k数据集。这是Databricks员工在2023年开源的数据集，专门用于监督微调，包含指令、上下文、响应和类别四个字段。分析了所有八个类别之后，我们决定只保留`information_extraction`类别中每条上下文的第一句话。

数据解析过程如下：

from datasets import load_dataset  
# 加载数据集
dataset = load_dataset("databricks/databricks-dolly-15k")
# 选择数据集中所需的类别
ie_category = [e for e in dataset["train"] if e["category"]=="information_extraction"]  
# 保留每个实例的上下文
ie_context = [e["context"] for e in ie_category]
# 将文本分割成句子（在句号处），并保留第一句话
reduced_context = [text.split('.')[0] + '.' for text in ie_context]
# 仅保留指定长度的序列（使用字符长度）
sampler = [e for e in reduced_context if 30 < len(e) < 170]

这样一通操作下来，我们拿到了1041个句子。因为这只是个迷你项目，所以没有手动逐个筛选句子，难免有些样本可能不太理想。如果是正式的生产项目，肯定要仔细挑一挑，不过对于这次实验来说，够用了。

格式化数据

首先创建一个系统消息，用来定义输入提示，并指示模型如何生成答案：

system_message = """您是一位经验丰富的注释者。  
从以下文本中提取所有实体及其之间的关系。 
将答案写成三元组实体1|关系|实体2。 
不要添加其他内容。 
示例文本：Alice is from France.  
答案：Alice|is from|France.  
"""

因为是实验阶段，提示词写得比较简单。当然也试过其他格式，包括要求以CoNLL格式输出（对实体进行分类），模型表现也还不错。但为了省事，这次就用最基础的方式。

接下来，把数据转换成对话格式：

messages = [[  
  {"role": "system","content": f"{system_message}"},  
  {"role": "user", "content": e}] for e in sampler]

Groq 客户端和 API

Llama3刚发布几天，API选项还比较有限。虽然Llama3-70B的聊天界面能用，但这个项目需要的是一个能批量处理1000个句子的API。GroqCloud的免费API正好派上了用场。

提醒一下：需要先登录GroqCloud网站，获取免费的API密钥，然后把密钥保存在Google Colab的secrets中。

初始化Groq客户端：

import os  
from groq import Groq  

gclient = Groq(  
  api_key=userdata.get("GROQ"),  
)

然后定义几个辅助函数，用来高效地与Meta.ai聊天界面交互：

import time              
from tqdm import tqdm              

def process_data(prompt):              
    """发送一个请求并检索模型的生成内容。"""              
    chat_completion = gclient.chat.completions.create(              
        messages=prompt, # 发送给模型的输入提示              
        model="llama3-70b-8192", # 根据 GroqCloud 标签              
        temperature=0.5, # 控制多样性              
        max_tokens=128, # 生成的最大标记数              
        top_p=1, # 考虑的可能性加权选项的比例              
        stop=None, # 表示停止生成的字符串              
        stream=False, # 如果设置，将发送部分消息              
    )              
    return chat_completion.choices[0].message.content              

def send_messages(messages):              
    """以批处理的方式处理消息，并在批处理之间暂停。"""              
    batch_size = 10  
    answers = []  
    for i in tqdm(range(0, len(messages), batch_size)): # 每批处理 10 条消息  
        batch = messages[i:i+10]  # 获取下一批消息  
        for message in batch:  
            output = process_data(message)  
            answers.append(output)  
        if i + 10 < len(messages):  # 检查是否还有批次  
            time.sleep(10)  # 等待 10 秒  
    return answers

process_data()是Groq客户端聊天完成函数的一个包装器。send_messages()则负责以小批量处理数据。为了避免超出免费API的请求和令牌限制，每处理10条消息后加了10秒的延迟。实际上不一定那么严格，但这样更保险。

现在，剩下的就很简单了：生成关系抽取数据，然后和原始数据集合并。

# 使用 Llama3-70B 进行数据生成
answers = send_messages(messages)

# 将输入数据与生成的数据集合并
combined_dataset = [{'text': user, 'gold_re': output} for user, output in zip(sampler, answers)]

评估 Llama3-8B 的基准性能

在正式进行微调之前，有必要先在几个样本上测试一下原版Llama3-8B的表现，看看微调是否真的有必要。

构建测试数据集

从刚构建的数据集中随机抽出20个样本作为测试集，剩下的都用来微调：

import random
random.seed(17)

# 选择 20 个随机条目
mini_data = random.sample(combined_dataset, 20)

# 构建对话格式
parsed_mini_data = [[{'role': 'system', 'content': system_message},
   {'role': 'user', 'content': e['text']}] for e in mini_data]

```python
{'text': 'Long before any knowledge of electricity existed, people were aware of shocks from electric fish.',
 'gold_re': 'people|were aware of|shocksnshocks|from|electric fishnelectric fish|had|electricity',
 'test_re': 'electric fish|were aware of|shocks'}

完整的测试数据集可以参考Google Colab笔记本。光从这一个样本就能看出来，Llama3-8B在关系抽取上确实还有不小的提升空间。那我们就想办法让它变强。

监督微调 Llama3-8B

微调过程中，我们会用到QLoRA和Flash Attention等一系列技术。关于超参数的选择，这里不展开细聊，感兴趣的话可以翻翻后面列出的参考资料。

A100 GPU支持Flash Attention和bfloat16，40GB左右的内存也足够应付这次微调。

准备 SFT 数据集

先把数据集解析成对话格式——系统消息、输入文本，以及从Llama3-70B生成中拿到的期望答案。然后保存为HuggingFace数据集：

def create_conversation(sample):
    return {
        "messages": [
            {"role": "system","content": system_message},
            {"role": "user", "content": sample["text"]},
            {"role": "assistant", "content": sample["gold_re"]}
        ]
    }

from datasets import load_dataset, Dataset

train_dataset = Dataset.from_list(train_data)

转换为对话格式：

train_dataset = train_dataset.map(create_conversation,
                                  remove_columns=train_dataset.features,
                                  batched=False)

选择模型：

model_id = "meta-llama/Meta-Llama-3-8B"

加载分词器：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_id,
                                          use_fast=True,
                                          trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
tokenizer.padding_side = 'left'

设置最大长度：

tokenizer.model_max_length = 512

选择量化参数：

from transformers import BitsAndBytesConfig  

bnb_config = BitsAndBytesConfig(  
    load_in_4bit=True,  
    bnb_4bit_use_double_quant=True,  
    bnb_4bit_quant_type="nf4",  
    bnb_4bit_compute_dtype=torch.bfloat16  
)

加载模型：

from transformers import AutoModelForCausalLM  
from peft import prepare_model_for_kbit_training  
from trl import setup_chat_format  

device_map = {"": torch.cuda.current_device()} if torch.cuda.is_a vailable() else None  

model = AutoModelForCausalLM.from_pretrained(  
    model_id,  
    device_map=device_map,  
    attn_implementation="flash_attention_2",  
    quantization_config=bnb_config  
)  

model, tokenizer = setup_chat_format(model, tokenizer)  
model = prepare_model_for_kbit_training(model)

LoRA 配置

from peft import LoraConfig  

# 根据Sebastian Raschka的研究结果  
peft_config = LoraConfig(  
    lora_alpha=128, #32  
    lora_dropout=0.05,  
    r=256,  #16  
    bias="none",  
    target_modules=["q_proj", "o_proj", "gate_proj", "up_proj",   
                    "down_proj", "k_proj", "v_proj"],  
    task_type="CAUSAL_LM",  
)

当目标覆盖所有线性层时，效果最好。如果内存有限，选更标准的值（比如alpha=32, rank=16）可以减少参数数量，也是不错的选择。

训练参数

from transformers import TrainingArguments

# 适配自 Phil Schmid 博客文章  
args = TrainingArguments(  
    output_dir=sft_model_path,  # 保存模型和存储库 ID 的目录  
    num_train_epochs=2,  # 训练周期数  
    per_device_train_batch_size=4,  # 训练期间每个设备的批处理大小  
    gradient_accumulation_steps=2,  # 执行向后/更新传递之前的步骤数  
    gradient_checkpointing=True,  # 使用梯度检查点以节省内存，在分布式训练中使用  
    optim="adamw_8bit",  # 如果内存不足，请选择 paged_adamw_8bit  
    logging_steps=10,  # 每 10 步记录一次日志  
    sa ve_strategy="epoch",  # 每个周期保存检查点  
    learning_rate=2e-4,  # 学习率，基于 QLoRA 论文  
    bf16=True,  # 使用 bfloat16 精度  
    tf32=True,  # 使用 tf32 精度  
    max_grad_norm=0.3,  # 基于 QLoRA 论文的最大梯度范数  
    warmup_ratio=0.03,  # 基于 QLoRA 论文的预热比例  
    lr_scheduler_type="constant",  # 使用恒定学习率调度程序  
    push_to_hub=True,   # 将模型推送到 Hugging Face hub  
    hub_model_id="llama3-8b-sft-qlora-re",  
    report_to="tensorboard",  # 报告指标到 tensorboard  
)

如果选择在本地保存模型，可以省略最后三个参数。另外，需要根据内存情况调整`per_device_batch_size`和`gradient_accumulation_steps`，防止出现OOM（内存不足）错误。

from trl import SFTTrainer  

trainer = SFTTrainer(  
    model=model,  
    args=args,  
    train_dataset=sft_dataset,  
    peft_config=peft_config,  
    max_seq_length=512,  
    tokenizer=tokenizer,  
    packing=False,  # True if the dataset is large  
    dataset_kwargs={  
        "add_special_tokens": False,  # the template adds the special tokens  
        "append_concat_token": False,  # no need to add additional separator token  
    }  
)  

trainer.train()  
trainer.sa ve_model()

训练和模型保存加起来，大约花了10分钟。

然后清理内存，为推理测试做准备。如果GPU内存较小，遇到了CUDA OOM错误，可能需要重启运行环境。

import torch  
import gc  
del model  
del tokenizer  
gc.collect()  
torch.cuda.empty_cache()

SFT 模型的推理测试

最后一步，加载半精度的基础模型和Peft适配器。这次测试，我们不将模型与适配器合并。

from peft import AutoPeftModelForCausalLM  
from transformers import AutoTokenizer, pipeline  
import torch

加载HF模型

peft_model_id = "solanaO/llama3-8b-sft-qlora-re"
# 使用 PEFT 适配器加载模型
model = AutoPeftModelForCausalLM.from_pretrained(
    peft_model_id,
    device_map="auto",
    torch_dtype=torch.float16,
    offload_buffers=True
)

加载分词器：

tokenizer = AutoTokenizer.from_pretrained(peft_model_id)

tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id

构建文本生成流水线：

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

加载测试数据集（包含之前选出的20个样本），以对话格式处理数据。这次，我们去掉助手消息部分，并格式化为HuggingFace数据集：

def create_input_prompt(sample):
    return {
        "messages": [
            {"role": "system","content": system_message},
            {"role": "user", "content": sample["text"]},
        ]
    }

from datasets import Dataset

test_dataset = Dataset.from_list(mini_data)