用开源小模型，
训练你自己的化学研究 Agent。

# pip install openai rdkit
# Generate tool-use trajectories with a teacher model and keep only the ones RDKit can verify.
import json, os
from openai import OpenAI
from rdkit import Chem
from rdkit.Chem import Descriptors, rdMolDescriptors

teacher = OpenAI(base_url="https://api.deepseek.com", api_key=os.environ["DEEPSEEK_API_KEY"])
TOOLS = [{
    "type": "function",
    "function": {
        "name": "describe_molecule",
        "description": "Validate a SMILES string and return its canonical form, formula, weight and heavy-atom count.",
        "parameters": {"type": "object", "properties": {"smiles": {"type": "string"}}, "required": ["smiles"]},
    },
}]

def describe_molecule(smiles: str) -> dict:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return {"valid": False}
    return {"valid": True, "canonical": Chem.MolToSmiles(mol), "formula": rdMolDescriptors.CalcMolFormula(mol),
            "mw": round(Descriptors.MolWt(mol), 2), "heavy_atoms": mol.GetNumHeavyAtoms()}

def run(task: str, max_steps: int = 6) -> list[dict]:
    messages = [{"role": "system", "content": "You are a careful chemist. Check every structure with the tool before answering."},
                {"role": "user", "content": task}]
    for _ in range(max_steps):
        msg = teacher.chat.completions.create(model="deepseek-flash", messages=messages, tools=TOOLS).choices[0].message
        turn = msg.model_dump(exclude_none=True)
        turn.pop("reasoning_content", None)          # keep the visible answer and tool calls only
        messages.append(turn)
        if not msg.tool_calls:
            return messages
        for call in msg.tool_calls:
            result = describe_molecule(**json.loads(call.function.arguments))
            messages.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result)})
    return messages

def verified(messages: list[dict]) -> bool:
    # Keep a trace only if the tool was used, every structure it saw was valid, and the run ended with an answer.
    results = [json.loads(m["content"]) for m in messages if m["role"] == "tool"]
    return bool(results) and all(r["valid"] for r in results) and messages[-1]["role"] == "assistant"

with open("traces.jsonl", "w", encoding="utf-8") as out:
    for task in open("tasks.txt", encoding="utf-8"):
        trace = run(task.strip())
        if verified(trace):
            out.write(json.dumps({"messages": trace, "tools": TOOLS}, ensure_ascii=False) + "\n")

# pip install "trl>=0.20" peft datasets accelerate
# LoRA fine-tuning on the verified traces. A 2–4B model fits on one 16–24 GB GPU.
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

BASE = "LiquidAI/LFM3-1.7B"   # or ibm-granite/granite-4.2-3b; for multimodal bases (Qwen3.5, Gemma 4) follow their model card
data = load_dataset("json", data_files="traces.jsonl", split="train").train_test_split(test_size=0.05, seed=0)

trainer = SFTTrainer(
    model=BASE,
    train_dataset=data["train"],
    eval_dataset=data["test"],
    peft_config=LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, target_modules="all-linear", task_type="CAUSAL_LM"),
    args=SFTConfig(
        output_dir="chem-agent-lora",
        num_train_epochs=2,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=8,
        learning_rate=1e-4,
        max_length=4096,
        gradient_checkpointing=True,
        bf16=True,
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=100,
    ),
)
trainer.train()
trainer.save_model("chem-agent-lora")

# Merge the adapter into the base weights so vLLM / llama.cpp can load a single model.
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer

model = AutoPeftModelForCausalLM.from_pretrained("chem-agent-lora")
model.merge_and_unload().save_pretrained("chem-agent-merged")
AutoTokenizer.from_pretrained("chem-agent-lora").save_pretrained("chem-agent-merged")

# pip install "trl>=0.20" rdkit datasets
# GRPO with a verifiable chemistry reward: RDKit grades every answer, no reward model needed.
import re
from datasets import Dataset
from rdkit import Chem, RDLogger
from trl import GRPOConfig, GRPOTrainer

RDLogger.DisableLog("rdApp.*")
ALLOWED = {"C", "H", "O", "F"}
TASK = ("Propose one lithium-battery electrolyte solvent containing only C, H, O and F, "
        "with at most 12 heavy atoms and at least two fluorine atoms. Answer as ....")
prompts = Dataset.from_list([{"prompt": [{"role": "user", "content": TASK}]}] * 256)   # use many varied tasks in practice

def text_of(completion):
    return completion[-1]["content"] if isinstance(completion, list) else completion

def chemistry_reward(completions, **kwargs):
    rewards = []
    for completion in completions:
        found = re.search(r"(.*?)", text_of(completion), re.S)
        mol = Chem.MolFromSmiles(found.group(1).strip()) if found else None
        if mol is None:
            rewards.append(-1.0)                      # no answer or an invalid structure
            continue
        symbols = [a.GetSymbol() for a in mol.GetAtoms()]
        score = 0.2                                   # a valid structure
        score += 0.4 if set(symbols) <= ALLOWED else -0.4
        score += 0.2 if mol.GetNumHeavyAtoms() <= 12 else -0.2
        score += 0.2 if symbols.count("F") >= 2 else 0.0
        rewards.append(score)
    return rewards

trainer = GRPOTrainer(
    model="chem-agent-merged",                        # start from the SFT model
    reward_funcs=chemistry_reward,
    train_dataset=prompts,
    args=GRPOConfig(output_dir="chem-agent-grpo", num_generations=8, max_completion_length=512,
                    per_device_train_batch_size=8, gradient_accumulation_steps=4,
                    learning_rate=5e-6, bf16=True, logging_steps=5),
)
trainer.train()
trainer.save_model("chem-agent-grpo")

# pip install mcp rdkit
from mcp.server.fastmcp import FastMCP
from rdkit import Chem
from rdkit.Chem import Descriptors, rdMolDescriptors

mcp = FastMCP("chem-tools")

@mcp.tool()
def describe_molecule(smiles: str) -> dict:
    """Validate a SMILES string and return canonical SMILES, formula, weight and heavy-atom count."""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return {"valid": False}
    return {"valid": True, "canonical": Chem.MolToSmiles(mol), "formula": rdMolDescriptors.CalcMolFormula(mol),
            "mw": round(Descriptors.MolWt(mol), 2), "heavy_atoms": mol.GetNumHeavyAtoms()}

if __name__ == "__main__":
    mcp.run()