docker run -itd --net=host \
    --name MolScribe \
    --shm-size=1g \
    --privileged \
    --device /dev/davinci0 \
    --device /dev/davinci1 \
    --device /dev/davinci2 \
    --device /dev/davinci3 \
    --device /dev/davinci4 \
    --device /dev/davinci5 \
    --device /dev/davinci6 \
    --device /dev/davinci7 \
    --device /dev/davinci8 \
    --device /dev/davinci9 \
    --device /dev/davinci10 \
    --device /dev/davinci11 \
    --device /dev/davinci12 \
    --device /dev/davinci13 \
    --device /dev/davinci14 \
    --device /dev/davinci15 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
    -v /usr/local/dcmi:/usr/local/dcmi \
    -v /usr/local/sbin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
    -v /weights/MolScribe:/weights/MolScribe \
    -v /home/tz/MolScribe:/home/tz/MolScribe \
    对应镜像/kubeflow/ascend-cann-mol:8.5.0-a3-ubuntu22.04-py3.11

docker exec -it MolScribe bash

export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
export PATH=/usr/local/python3.11.14/bin:$PATH
export PYTHONPATH=/home/tz/MolScribe:$PYTHONPATH

export ASCEND_RT_VISIBLE_DEVICES=0
export ASCEND_OPP_PATH=/usr/local/Ascend/cann-8.5.0/opp
export ASCEND_AICPU_PATH=/usr/local/Ascend/cann-8.5.0
export ASCEND_HOME_PATH=/usr/local/Ascend/cann-8.5.0
export ASCEND_TOOLKIT_HOME=/usr/local/Ascend/cann-8.5.0
export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/lib64:/usr/local/Ascend/driver/lib64:$LD_LIBRARY_PATH

# 安装 PyTorch NPU
pip3 install torch==2.5.1 torch_npu==2.5.1 torchvision==0.20.1 -f https://download.pytorch.org/whl/torch_npu.html

# 安装其他依赖
pip3 install "numpy>=1.19.5,<2.0"
pip3 install pandas>=1.2.4 matplotlib>=3.5.3
pip3 install opencv-python-headless==4.8.1.78
pip3 install transformers>=4.5.1 huggingface-hub>=0.11.0 tensorboardX timm==0.4.12
pip3 install rdkit>=2022.3.3 albumentations==1.1.0
pip3 install PyYAML>=5.4.1 scikit-learn>=1.0.0 scipy>=1.7.0
pip3 install SmilesPE==0.0.3 "OpenNMT-py==2.2.0" Pillow>=9.0.0
pip3 install scikit-image>=0.16.1 qudida>=0.0.4 "pyonmttok>=1.23,<2"
pip3 install decorator attrs psutil

cd /home/tz/MolScribe
mkdir -p /weights/MolScribe/molscribe_output
mkdir -p /weights/MolScribe/ckpts

# 使用 Hugging Face 镜像站下载
export HF_ENDPOINT=https://hf-mirror.com

python3 -c "
from huggingface_hub import hf_hub_download
model_path = hf_hub_download(
    repo_id='yujieq/MolScribe',
    filename='swin_base_char_aux_1m.pth',
    repo_type='model'
)
print('✅ 模型下载成功:', model_path)
"

git apply all-npu-fix.patch

cat > predict.py << 'EOF'
import argparse
import json

# rdkit 必须在 torch 之前导入，以避免段错误
import rdkit
import rdkit.Chem as Chem
import torch

# NPU 支持：导入 torch_npu
# 注意：transfer_to_npu 会自动patch PyTorch，可能导致CPU模式问题
try:
    import torch_npu
except ImportError:
    pass

from molscribe import MolScribe

import warnings
warnings.filterwarnings('ignore')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str, default=None, required=True)
    parser.add_argument('--image_path', type=str, default=None, required=True)
    parser.add_argument('--return_confidence', action='store_true')
    parser.add_argument('--return_atoms_bonds', action='store_true')
    parser.add_argument('--device', type=str, default='npu', choices=['npu', 'cuda', 'cpu'])
    args = parser.parse_args()

    # 设备选择：优先使用 NPU，其次 CUDA，最后 CPU
    if args.device == 'npu':
        if torch.npu.is_available():
            device = torch.device('npu')
            print(f"Using NPU device: {torch.npu.get_device_name(0)}")
        elif torch.cuda.is_available():
            device = torch.device('cuda')
            print(f"NPU not available, falling back to CUDA: {torch.cuda.get_device_name(0)}")
        else:
            device = torch.device('cpu')
            print("NPU and CUDA not available, using CPU")
    elif args.device == 'cuda':
        if torch.cuda.is_available():
            device = torch.device('cuda')
            print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
        else:
            device = torch.device('cpu')
            print("CUDA not available, using CPU")
    else:
        device = torch.device('cpu')
        print("Using CPU device")

    model = MolScribe(args.model_path, device)
    output = model.predict_image_file(
        args.image_path, return_atoms_bonds=args.return_atoms_bonds, return_confidence=args.return_confidence)
    for key, value in output.items():
        print(f"{key}:")
        print(value + '\n' if isinstance(value, str) else json.dumps(value) + '\n')
EOF

# 使用 Python 脚本自动修复
python3 << 'PYEOF'
with open('molscribe/interface.py', 'r') as f:
    content = f.read()

# 替换导入顺序
old = '''import cv2
import torch'''
new = '''import cv2
# rdkit 必须在 torch 之前导入，避免段错误
import rdkit
import rdkit.Chem as Chem
import torch'''

content = content.replace(old, new)

# 在 torch 之后添加 torch_npu 导入
old2 = '''import torch
import numpy as np'''
new2 = '''import torch
import numpy as np

# NPU 支持
try:
    import torch_npu
except ImportError:
    pass'''

content = content.replace(old2, new2)

with open('molscribe/interface.py', 'w') as f:
    f.write(content)

print("✅ molscribe/interface.py 修复完成")
PYEOF

python3 << 'PYEOF'
import rdkit
import rdkit.Chem as Chem
import torch
import json
from molscribe import MolScribe

device = torch.device('npu')
print("🚀 正在加载模型（NPU）...")

model_path = "/weights/MolScribe/ckpts/swin_base_char_aux_1m.pth"
model = MolScribe(model_path, device=device)

print("✅ 模型加载成功！")

print("🧪 运行推理...")
output = model.predict_image_file(
    "assets/example.png",
    return_atoms_bonds=True,
    return_confidence=True
)

print("🔍 SMILES:", output["smiles"])
print("📊 Confidence:", output["confidence"])

with open("/weights/MolScribe/molscribe_output/result.json", "w") as f:
    json.dump(output, f, indent=2, default=str)

print("✅ 推理结果已保存至：/weights/MolScribe/molscribe_output/result.json")
PYEOF

cd /home/tz/MolScribe

export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_DISABLE_PROXY=1

python3 -c "
from huggingface_hub import hf_hub_download
real_zip = hf_hub_download(
    repo_id='yujieq/MolScribe',
    filename='real.zip',
    repo_type='model'
)
print('✅ 下载成功:', real_zip)
"

python3 -c "
import zipfile
import glob
paths = glob.glob('/root/.cache/huggingface/hub/models--yujieq--MolScribe/**/real.zip', recursive=True)
real_zip = paths[0]
target_dir = '/home/tz/MolScribe/data'
import os
os.makedirs(target_dir, exist_ok=True)
with zipfile.ZipFile(real_zip, 'r') as zf:
    zf.extractall(target_dir)
print('✅ 解压完成:', target_dir)
"

# 保存为 /home/tz/MolScribe/batch_predict.py
"""
批量推理脚本 - 在 NPU 上运行 MolScribe 推理
"""
import os
import csv
import json
import pandas as pd

# 环境变量设置（必须在导入 torch 前）
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
os.environ['ASCEND_OPP_PATH'] = '/usr/local/Ascend/cann-8.5.0/opp'
os.environ['ASCEND_AICPU_PATH'] = '/usr/local/Ascend/cann-8.5.0'
os.environ['ASCEND_HOME_PATH'] = '/usr/local/Ascend/cann-8.5.0'
os.environ['ASCEND_TOOLKIT_HOME'] = '/usr/local/Ascend/cann-8.5.0'
os.environ['LD_LIBRARY_PATH'] = '/usr/local/Ascend/cann-8.5.0/lib64:/usr/local/Ascend/driver/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')

# rdkit 必须在 torch 之前导入
import rdkit
import rdkit.Chem as Chem
import torch
try:
    import torch_npu
except ImportError:
    pass

from molscribe import MolScribe


def batch_predict(model, image_files, batch_size=8):
    all_predictions = []
    for i in range(0, len(image_files), batch_size):
        batch_files = image_files[i:i+batch_size]
        print(f"Processing {i+1}-{min(i+batch_size, len(image_files))}/{len(image_files)}")
        results = model.predict_image_files(batch_files, return_confidence=True)
        all_predictions.extend(results)
    return all_predictions


def main():
    data_dir = '/home/tz/MolScribe/data'
    model_path = '/weights/MolScribe/ckpts/swin_base_char_aux_1m.pth'
    output_dir = '/weights/MolScribe/molscribe_output'
    os.makedirs(output_dir, exist_ok=True)

    datasets = ["acs", "CLEF", "staker", "UOB", "USPTO"]
    for ds in datasets:
        print(f"\n🚀 测试数据集: {ds}")
        gold_file = os.path.join(data_dir, f'real/{ds}.csv')

        gold_df = pd.read_csv(gold_file)
        image_files = []
        image_ids = []
        for _, row in gold_df.iterrows():
            img_path = os.path.join(data_dir, row['file_path'])
            if os.path.exists(img_path):
                image_files.append(img_path)
                image_ids.append(row['image_id'])
            else:
                print(f"⚠️ 图像缺失: {img_path}")

        print(f"✅ 共处理 {len(image_files)} 张图片")

        print("🧠 加载模型...")
        device = torch.device('npu')
        model = MolScribe(model_path, device=device)

        print("🚀 批量推理...")
        predictions = batch_predict(model, image_files, batch_size=8)

        # 保存结果
        pred_csv = os.path.join(output_dir, f"prediction_{ds}.csv")
        with open(pred_csv, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['image_id', 'SMILES'])
            for img_id, pred in zip(image_ids, predictions):
                writer.writerow([img_id, pred['smiles']])

        detailed = [{'image_id': img_id, 'smiles': pred['smiles'], 'confidence': pred.get('confidence')} 
                   for img_id, pred in zip(image_ids, predictions)]
        with open(os.path.join(output_dir, f"prediction_{ds}_detailed.json"), 'w') as f:
            json.dump(detailed, f, indent=2)

        print(f"✅ 结果已保存: {pred_csv}")
    print("\n🎉 所有数据集测试完成！")


if __name__ == '__main__':
    main()

cd /home/tz/MolScribe
python3 batch_predict.py

cd /home/tz/MolScribe && python3 << 'PYEOF'
import os
import subprocess
import json

datasets = ["acs", "CLEF", "staker", "UOB", "USPTO"]
data_dir = "/home/tz/MolScribe/data/real"
pred_dir = "/weights/MolScribe/molscribe_output"
result_md = "/home/tz/MolScribe/测试结果.md"

results = {}

print("=" * 60)
print(" 开始全自动测试所有数据集：".center(60))
print("=" * 60)

for ds in datasets:
    print(f"\n🚀 正在测试：{ds}")
    cmd = f'sed -i "s/dataset_name = .*/dataset_name = \\"{ds}\\"/" batch_predict.py'
    os.system(cmd)
    subprocess.run(["python3", "batch_predict.py"], check=True)

    gold_csv = os.path.join(data_dir, f"{ds}.csv")
    pred_csv = os.path.join(pred_dir, f"prediction_{ds}.csv")
    res = subprocess.run(
        ["python3", "evaluate.py", "--gold_file", gold_csv, "--pred_file", pred_csv, "--pred_field", "SMILES"],
        capture_output=True, text=True
    )

    try:
        output = res.stdout.strip()
        data = json.loads(output[output.find("{"):output.rfind("}")+1])
        results[ds] = data
        print(f"✅ {ds} 成功：{data}")
    except Exception as e:
        print(f"❌ {ds} 失败: {e}")
        results[ds] = None

# 生成 Markdown 报告
with open(result_md, "w", encoding="utf-8") as f:
    f.write("# MolScribe 昇腾 NPU 精度测试报告\n\n")
    f.write("| 数据集 | canon_smiles | graph | chiral |\n")
    f.write("|--------|-------------|-------|--------|\n")
    for name, res in results.items():
        if res:
            ca = f"{res['canon_smiles']:.4f}"
            g = f"{res['graph']:.4f}"
            c = f"{res['chiral']:.4f}"
            f.write(f"| {name} | {ca} | {g} | {c} |\n")
        else:
            f.write(f"| {name} | 失败 | 失败 | 失败 |\n")

print("\n" + "=" * 60)
print(f"✅ 全部测试完成！报告已保存至：\n{result_md}")
print("=" * 60)
PYEOF

#!/bin/bash
# deploy_molscribe.sh - 一键部署 MolScribe 于昇腾 NPU 环境

set -e

echo "🔄 正在部署 MolScribe 于昇腾 NPU..."

# 1. 创建容器
docker rm -f MolScribe || true
docker run -itd --net=host \
    --name MolScribe \
    --shm-size=1g \
    --privileged \
    --device /dev/davinci0 \
    --device /dev/davinci1 \
    --device /dev/davinci2 \
    --device /dev/davinci3 \
    --device /dev/davinci4 \
    --device /dev/davinci5 \
    --device /dev/davinci6 \
    --device /dev/davinci7 \
    --device /dev/davinci8 \
    --device /dev/davinci9 \
    --device /dev/davinci10 \
    --device /dev/davinci11 \
    --device /dev/davinci12 \
    --device /dev/davinci13 \
    --device /dev/davinci14 \
    --device /dev/davinci15 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
    -v /usr/local/dcmi:/usr/local/dcmi \
    -v /usr/local/sbin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
    -v /weights/MolScribe:/weights/MolScribe \
    -v /home/tz/MolScribe:/home/tz/MolScribe \
    harbor.stonewise.cn/kubeflow/ascend-cann-mol:8.5.0-a3-ubuntu22.04-py3.11

# 2. 进入容器并执行部署
docker exec -it MolScribe bash << 'EOF'

# 设置路径与环境
export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
export PATH=/usr/local/python3.11.14/bin:$PATH
export PYTHONPATH=/home/tz/MolScribe:$PYTHONPATH

export ASCEND_RT_VISIBLE_DEVICES=0
export ASCEND_OPP_PATH=/usr/local/Ascend/cann-8.5.0/opp
export ASCEND_AICPU_PATH=/usr/local/Ascend/cann-8.5.0
export ASCEND_HOME_PATH=/usr/local/Ascend/cann-8.5.0
export ASCEND_TOOLKIT_HOME=/usr/local/Ascend/cann-8.5.0
export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/lib64:/usr/local/Ascend/driver/lib64:$LD_LIBRARY_PATH

cd /home/tz/MolScribe
mkdir -p /weights/MolScribe/molscribe_output /weights/MolScribe/ckpts

# 安装依赖
pip3 install torch==2.5.1 torch_npu==2.5.1 torchvision==0.20.1 -f https://download.pytorch.org/whl/torch_npu.html
pip3 install "numpy>=1.19.5,<2.0" pandas matplotlib opencv-python-headless==4.8.1.78 transformers huggingface-hub tensorboardX timm==0.4.12 rdkit albumentations PyYAML scikit-learn scipy SmilesPE OpenNMT-py scikit-image qudida pyonmttok decorator attrs psutil

# 下载模型
export HF_ENDPOINT=https://hf-mirror.com
python3 -c "
from huggingface_hub import hf_hub_download
model_path = hf_hub_download(
    repo_id='yujieq/MolScribe',
    filename='swin_base_char_aux_1m.pth',
    repo_type='model'
)
print('✅ 模型下载成功:', model_path)
"

# 修复代码
python3 << 'PYEOF'
with open('molscribe/interface.py', 'r') as f:
    content = f.read()
old = '''import cv2
import torch'''
new = '''import cv2
# rdkit 必须在 torch 之前导入，避免段错误
import rdkit
import rdkit.Chem as Chem
import torch'''
content = content.replace(old, new)
old2 = '''import torch
import numpy as np'''
new2 = '''import torch
import numpy as np

# NPU 支持
try:
    import torch_npu
except ImportError:
    pass'''
content = content.replace(old2, new2)
with open('molscribe/interface.py', 'w') as f:
    f.write(content)
print("✅ 代码修复完成")
PYEOF

# 运行单图测试
python3 << 'PYEOF'
import rdkit
import rdkit.Chem as Chem
import torch
import json
from molscribe import MolScribe

device = torch.device('npu')
print("🚀 正在加载模型（NPU）...")
model = MolScribe("/weights/MolScribe/ckpts/swin_base_char_aux_1m.pth", device=device)
print("✅ 模型加载成功！")

output = model.predict_image_file("assets/example.png", return_atoms_bonds=True, return_confidence=True)
print("SMILES:", output["smiles"])
print("Confidence:", output["confidence"])
with open("/weights/MolScribe/molscribe_output/result.json", "w") as f:
    json.dump(output, f, indent=2, default=str)
print("✅ 推理结果已保存")
PYEOF

# 生成批量推理脚本
cat > batch_predict.py << 'EOF'
# 请粘贴上文中的 batch_predict.py 内容
EOF

# 下载并解压数据集
export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_DISABLE_PROXY=1
python3 -c "
from huggingface_hub import hf_hub_download
real_zip = hf_hub_download(
    repo_id='yujieq/MolScribe',
    filename='real.zip',
    repo_type='model'
)
print('✅ 下载成功:', real_zip)
"
python3 -c "
import zipfile
import glob
paths = glob.glob('/root/.cache/huggingface/hub/models--yujieq--MolScribe/**/real.zip', recursive=True)
real_zip = paths[0]
target_dir = '/home/tz/MolScribe/data'
import os
os.makedirs(target_dir, exist_ok=True)
with zipfile.ZipFile(real_zip, 'r') as zf:
    zf.extractall(target_dir)
print('✅ 解压完成:', target_dir)
"

# 运行批量测试
python3 batch_predict.py

# 生成报告
cd /home/tz/MolScribe && python3 << 'PYEOF'
import os
import subprocess
import json
datasets = ["acs", "CLEF", "staker", "UOB", "USPTO"]
data_dir = "/home/tz/MolScribe/data/real"
pred_dir = "/weights/MolScribe/molscribe_output"
result_md = "/home/tz/MolScribe/测试结果.md"
results = {}
for ds in datasets:
    os.system(f'sed -i "s/dataset_name = .*/dataset_name = \\"{ds}\\"/" batch_predict.py')
    subprocess.run(["python3", "batch_predict.py"], check=True)
    gold_csv = os.path.join(data_dir, f"{ds}.csv")
    pred_csv = os.path.join(pred_dir, f"prediction_{ds}.csv")
    res = subprocess.run(
        ["python3", "evaluate.py", "--gold_file", gold_csv, "--pred_file", pred_csv, "--pred_field", "SMILES"],
        capture_output=True, text=True
    )
    try:
        data = json.loads(res.stdout.strip()[res.stdout.strip().find("{"):res.stdout.strip().rfind("}")+1])
        results[ds] = data
    except:
        results[ds] = None
with open(result_md, "w", encoding="utf-8") as f:
    f.write("# MolScribe 昇腾 NPU 精度测试报告\n\n")
    f.write("| 数据集 | canon_smiles | graph | chiral |\n")
    f.write("|--------|-------------|-------|--------|\n")
    for name, res in results.items():
        ca = f"{res['canon_smiles']:.4f}" if res else "失败"
        g = f"{res['graph']:.4f}" if res else "失败"
        c = f"{res['chiral']:.4f}" if res else "失败"
        f.write(f"| {name} | {ca} | {g} | {c} |\n")
print("🎉 部署与测试完成！")
PYEOF

EOF

echo "✅ MolScribe 部署与测试已完成！"
echo "📄 报告路径：/home/tz/MolScribe/测试结果.md"