langchain和pytorch结合笔记

文章目录

- - 代码

只是一个可以跑通的示例，主内容见其他文章。

代码

python 复制代码

import torch
import gc
import re
import json
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_core.tools import tool
from langchain_huggingface import HuggingFacePipeline
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# --- 1. 环境清理 ---
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

# --- 2. 定义 PyTorch 工具 ---
from transformers import pipeline as hf_pipeline

print("加载工具 (CPU)...")
classifier = hf_pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english", device=-1)


@tool
def sentiment_analyzer(text: str) -> str:
    """分析文本的情感倾向 (POSITIVE/NEGATIVE)"""
    if not text: return "输入为空"
    result = classifier(text)
    return f"情感: {result[0]['label']}, 置信度: {result[0]['score']:.4f}"


# --- 3. 加载 Qwen2.5-0.5B ---
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
print("加载 Qwen2.5-0.5B...")

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)
model = model.eval()

# --- 4. 构建 Pipeline ---
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.3,  # 低温度，让模型更确定地输出工具格式
    top_p=0.9,
    pad_token_id=tokenizer.eos_token_id,
    return_full_text=False
)

llm = HuggingFacePipeline(pipeline=pipe)

# --- 5. 手动构建"Agent"逻辑 (LCEL) ---

# 定义工具调用的提示词模板
# 我们教模型：如果要调用工具，必须用特定的 XML 标签包裹


system_prompt = """你是一个智能助手。你可以使用以下工具：
1. sentiment_analyzer: 分析文本情感。
如果你需要使用工具，请严格按照以下格式回答：
<|tool_call_start|>sentiment_analyzer<|tool_call_end|>{{"text": "需要分析的文本"}}<|tool_call_end|>
如果不需要工具，直接回答用户的问题。
"""

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}")
])


# 定义一个简单的"执行器"函数
def run_agent_step(inputs):
    user_input = inputs["input"]

    # 1. 调用 LLM
    response = (prompt | llm | StrOutputParser()).invoke({"input": user_input})
    print(f"🤖 模型原始输出: {response}")

    # 2. 检查是否包含工具调用标签
    if "<|tool_call_start|>" in response:
        # 提取工具参数 (简单的正则提取)
        try:
            # 提取 {"text": "..."} 部分
            json_str = re.search(r'\{.*\}', response, re.DOTALL).group()
            tool_input = json.loads(json_str)

            print(f"🛠️ 正在调用工具: sentiment_analyzer")
            # 3. 执行工具
            tool_result = sentiment_analyzer.invoke(tool_input)

            return f"工具返回结果：{tool_result}。根据这个结果，用户的输入表达了{tool_result.split(',')[0].split(': ')[1]}的情绪。"
        except Exception as e:
            return f"工具调用失败: {e}。原始输出: {response}"
    else:
        return response


# --- 6. 运行测试 ---
print("\n--- 开始测试 ---")
try:
    # 直接调用我们定义的函数，而不是复杂的 agent.invoke
    result = run_agent_step({"input": "分析一下 'I love this phone' 的情绪"})
    print(f"\n✅ 最终回答: {result}")
except Exception as e:
    print(f"❌ 发生错误: {e}")
    import traceback

    traceback.print_exc()

输出结果：

bash 复制代码

加载工具 (CPU)...
加载 Qwen2.5-0.5B...

--- 开始测试 ---
🤖 模型原始输出: 。

Assistant: <|tool_call_start|>sentiment_analyzer<|tool_call_end|>{"text": "I love this phone"}<|tool_call_end|>
🛠️ 正在调用工具: sentiment_analyzer
✅ 最终回答: 工具返回结果：情感: POSITIVE, 置信度: 0.9998。根据这个结果，用户的输入表达了POSITIVE的情绪。