使用 Hugging Face Transformers 和 PyTorch 實現信息抽取
在自然語言處理(NLP)領域,信息抽取是一種常見的任務,其目標是從文本中提取特定類型的結構化信息。本文將介紹如何使用 Hugging Face Transformers 和 PyTorch 實現基于大語言模型的信息抽取。我們將通過一個具體的例子,展示如何從文本中抽取商品的詳細信息,包括產品名稱、品牌、特點、價格和銷量等。
1. 任務背景
信息抽取在許多應用場景中都非常有用,例如從新聞文章中提取關鍵信息、從產品描述中提取商品信息等。傳統的信息抽取方法通常依賴于規則匹配或基于機器學習的分類器,但這些方法往往需要大量的標注數據,并且在面對復雜的文本結構時效果有限。近年來,隨著大語言模型(如 GPT、Bert 等)的發展,基于這些模型的信息抽取方法逐漸成為研究熱點。
2. 技術棧介紹
2.1 Hugging Face Transformers
Hugging Face Transformers 是一個開源的 Python 庫,提供了大量預訓練的語言模型,以及用于微調和部署這些模型的工具。它支持多種模型架構(如 Bert、GPT、T5 等),并且提供了豐富的 API,方便開發者快速實現各種 NLP 任務。
2.2 PyTorch
PyTorch 是一個流行的深度學習框架,以其動態計算圖和易用性而聞名。它提供了強大的 GPU 支持,能夠加速模型的訓練和推理過程。在本文中,我們將使用 PyTorch 來加載和運行預訓練的語言模型。
3. 實現步驟
3.1 定義任務和數據
我們的目標是從給定的文本中抽取商品的相關信息。具體來說,我們需要提取以下屬性:
- 產品名稱
- 品牌
- 特點
- 原價
- 促銷價
- 銷量
為了實現這一目標,我們定義了一個簡單的數據結構來描述這些屬性,并提供了一些示例數據供模型學習。
3.2 初始化模型和分詞器
我們使用 Hugging Face Transformers 提供的 AutoTokenizer
和 AutoModelForCausalLM
來加載預訓練的語言模型。在這個例子中,我們使用了 Qwen2.5-1.5B-Instruct
模型,這是一個經過指令微調的模型,能夠更好地理解自然語言指令。
tokenizer = AutoTokenizer.from_pretrained(r"C:\Users\妄生\Qwen2.5-1.5B-Instruct", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(r"C:\Users\妄生\Qwen2.5-1.5B-Instruct", trust_remote_code=True)
3.3 構建提示模板
為了幫助模型更好地理解任務,我們設計了一個提示模板,明確告訴模型需要抽取哪些信息。模板中包含了需要抽取的實體類型和屬性,以及如何格式化輸出結果。
IE_PATTERN = "{}\n\n提取上述句子中{}的實體,并按照JSON格式輸出,上述句子中不存在的信息用['原文中未提及']來表示,多個值之間用','分隔。"
3.4 準備示例數據
為了讓模型更好地理解任務,我們提供了一些示例數據。這些示例數據展示了如何從文本中抽取信息,并以 JSON 格式輸出。
ie_examples = {'商品': [{'content': '2024 年新款時尚運動鞋,品牌 JKL,舒適透氣,多種顏色可選。原價 599 元,現在促銷價 499 元。月銷量 2000 雙。','answers': {'產品': ['時尚運動鞋'],'品牌': ['JKL'],'特點': ['舒適透氣、多種顏色可選'],'原價': ['599元'],'促銷價': ['499元'],'月銷量': ['2000雙'],}}]
}
3.5 實現信息抽取函數
我們定義了一個 inference
函數,用于對輸入的文本進行信息抽取。這個函數首先將輸入文本和提示模板組合起來,然后使用模型生成輸出結果。最后,我們將生成的結果解碼并打印出來。
def inference(sentences: list, custom_settings: dict):for sentence in sentences:cls_res = "商品"if cls_res not in schema:print(f'The type model inferenced {cls_res} which is not in schema dict, exited.')exit()properties_str = ', '.join(schema[cls_res])schema_str_list = f'"{cls_res}"({properties_str})'sentence_with_ie_prompt = IE_PATTERN.format(sentence, schema_str_list)full_input_text = build_prompt(sentence_with_ie_prompt, custom_settings["ie_pre_history"])inputs = tokenizer(full_input_text, return_tensors="pt")inputs = inputs.to('cpu')output_sequences = model.generate(inputs["input_ids"], max_length=2000, attention_mask=inputs.attention_mask)generated_reply = tokenizer.decode(output_sequences[0], skip_special_tokens=True)print(generated_reply[len(full_input_text):])
3.6 測試模型
最后,我們使用一些測試數據來驗證模型的效果。
sentences = ['“2024 年潮流雙肩包,品牌 PQR,材質耐磨,內部空間大。定價399元,優惠后價格349元。周銷量800個。”','2024 年智能手表,品牌為華為,功能強大,續航持久。售價1299元,便宜價999元,目前有黑色和銀色可選。月銷量3000只。'
]
custom_settings = init_prompts()
inference(sentences, custom_settings)
4. 結果展示
運行上述代碼后,模型將輸出從測試數據中抽取的信息。例如,對于第一句測試數據,模型可能會輸出類似以下的結果:
5. 總結
本文介紹了一個基于 Hugging Face Transformers 和 PyTorch 的信息抽取實現。通過使用預訓練的語言模型和提示模板,我們能夠從文本中抽取特定的結構化信息。這種方法的優點是不需要大量的標注數據,并且能夠處理復雜的文本結構。然而,它也有一些局限性,例如對模型的依賴較大,且生成的結果可能需要進一步驗證和修正。
未來,我們可以探索更多優化方法,例如微調模型以提高其在特定任務上的性能,或者結合其他技術(如規則匹配)來提高抽取的準確性。此外,我們還可以嘗試將這種方法應用于其他類型的信息抽取任務,以驗證其通用性。