MGeo实战技巧：如何修改推理.py脚本自定义输入输出格式

1. 背景与应用场景

在实体对齐任务中，地址数据的标准化和相似度匹配是关键环节。阿里开源的MGeo模型专注于中文地址领域的语义理解与相似度计算，能够高效识别不同表述但指向同一地理位置的地址对。该模型已在多个城市服务、物流调度和地图平台中落地应用。

默认情况下，推理.py脚本采用固定的输入输出格式，通常为 JSON 或 CSV 文件读取，并以标准结构返回相似度分数。然而，在实际工程场景中，业务系统往往需要适配特定的数据源格式（如数据库直连、API 流式输入）或定制化输出结构（如嵌入到推荐系统、风控引擎）。因此，掌握如何自定义输入输出格式成为提升 MGeo 集成效率的核心技能。

本文将基于已部署的 MGeo 镜像环境，手把手教你修改推理.py脚本，实现灵活的数据接口适配，满足多样化生产需求。

2. 环境准备与基础运行

2.1 镜像部署与环境激活

MGeo 支持通过容器化镜像快速部署，适用于单卡 GPU 环境（如 4090D），具体步骤如下：

# 启动镜像（示例命令） docker run -it --gpus all -p 8888:8888 mgeo-chinese-address:v1.0 # 进入容器后激活 Conda 环境 conda activate py37testmaas

2.2 Jupyter 交互开发支持

启动 Jupyter Notebook 可实现可视化调试：

jupyter notebook --ip=0.0.0.0 --port=8888 --allow-root

建议将核心脚本复制至工作区以便编辑：

cp /root/推理.py /root/workspace

随后可在/root/workspace目录下使用文本编辑器或 Jupyter Lab 打开并修改推理.py。

3. 修改推理脚本实现自定义输入输出

3.1 原始脚本结构分析

原始推理.py一般包含以下模块：

• 模型加载逻辑
• 输入数据读取（常为input.json）
• 地址对预处理
• 推理执行
• 输出结果写入文件（如output.json）

典型调用方式：

data = json.load(open("input.json", "r", encoding="utf-8")) # ... 模型推理 ... json.dump(results, open("output.json", "w", encoding="utf-8"), ensure_ascii=False)

这种硬编码路径和固定格式不利于集成。我们需将其改造为可配置、可扩展的形式。

3.2 自定义输入格式：支持多种数据源

支持命令行参数传入输入路径

引入argparse模块，允许用户指定输入文件：

import argparse def parse_args(): parser = argparse.ArgumentParser(description="MGeo Address Matching Inference") parser.add_argument("--input_path", type=str, default="input.json", help="Path to input data file") parser.add_argument("--output_path", type=str, default="output.json", help="Path to save results") parser.add_argument("--format", type=str, choices=["json", "csv"], default="json", help="Input/output format") return parser.parse_args()

动态解析不同格式输入

根据--format参数自动选择解析方式：

import pandas as pd import json def load_data(file_path, fmt): if fmt == "json": with open(file_path, "r", encoding="utf-8") as f: return json.load(f) elif fmt == "csv": df = pd.read_csv(file_path) # 假设CSV有两列：addr1, addr2 return df[["addr1", "addr2"]].to_dict(orient="records") else: raise ValueError(f"Unsupported format: {fmt}")

示例输入兼容性增强

支持更灵活的字段命名映射：

def normalize_input(data): normalized = [] for item in data: if isinstance(item, dict): addr1 = item.get("address1") or item.get("addr1") or item.get("source_addr") addr2 = item.get("address2") or item.get("addr2") or item.get("target_addr") pair_id = item.get("id", len(normalized)) normalized.append({"id": pair_id, "addr1": addr1, "addr2": addr2}) return normalized

3.3 自定义输出格式：结构化与可扩展输出

多格式输出支持

类似地，添加输出格式控制逻辑：

def save_results(results, output_path, fmt): if fmt == "json": with open(output_path, "w", encoding="utf-8") as f: json.dump(results, f, ensure_ascii=False, indent=2) elif fmt == "csv": df = pd.DataFrame(results) df.to_csv(output_path, index=False, encoding="utf_8_sig") # 兼容Excel

输出字段增强建议

除默认相似度分数外，建议增加以下元信息便于下游使用：

• match_score: 相似度得分（0~1）
• is_match: 判定是否为同一地点（阈值可配置）
• inference_time: 单次推理耗时（ms）
• model_version: 当前模型版本号

result_entry = { "id": pair["id"], "addr1": pair["addr1"], "addr2": pair["addr2"], "match_score": float(similarity), "is_match": bool(similarity > threshold), "inference_time_ms": round(time_cost * 1000, 2), "model_version": "mgeo-v1.2-zh-address" }

3.4 实际修改后的完整流程示例

以下是整合后的主函数框架：

import time def main(): args = parse_args() # 加载模型（假设已有 model 对象） from mgeo_model import MGeoModel model = MGeoModel.load_from_checkpoint("/root/model.ckpt") model.eval() # 加载并标准化输入 raw_data = load_data(args.input_path, args.format) inputs = normalize_input(raw_data) results = [] threshold = 0.85 # 可通过参数传入 for pair in inputs: start_t = time.time() try: similarity = model.predict(pair["addr1"], pair["addr2"]) time_cost = time.time() - start_t result_entry = { "id": pair["id"], "addr1": pair["addr1"], "addr2": pair["addr2"], "match_score": float(similarity), "is_match": bool(similarity > threshold), "inference_time_ms": round(time_cost * 1000, 2), "status": "success" } except Exception as e: result_entry = { "id": pair["id"], "addr1": pair["addr1"], "addr2": pair["addr2"], "match_score": 0.0, "is_match": False, "inference_time_ms": 0, "status": f"error: {str(e)}" } results.append(result_entry) # 保存结果 save_results(results, args.output_path, args.format) print(f"Inference completed. Results saved to {args.output_path}")

调用方式变为：

python 推理.py --input_path /data/test.csv --output_path /result/matches.json --format csv

4. 工程优化与最佳实践

4.1 批量推理性能优化

对于大规模地址对匹配任务，应启用批量处理模式：

# 修改模型预测接口支持 batch def batch_predict(self, addr1_list, addr2_list): embeddings1 = self.encode(addr1_list) embeddings2 = self.encode(addr2_list) similarities = cosine_similarity(embeddings1, embeddings2) return similarities.diagonal().tolist()

配合分批读取，减少 I/O 压力：

from itertools import islice def batch_iter(data, batch_size=32): it = iter(data) while True: batch = list(islice(it, batch_size)) if not batch: break yield batch

4.2 错误处理与日志记录

添加异常捕获和日志输出，提高稳定性：

import logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[logging.StreamHandler()] ) try: similarity = model.predict(addr1, addr2) except RuntimeError as e: logging.warning(f"Prediction failed for ID {pair['id']}: {e}") similarity = 0.0

4.3 配置文件替代硬编码参数

建议将阈值、模型路径等抽象为配置文件（如config.yaml）：

model: path: /root/model.ckpt threshold: 0.85 max_length: 128 io: default_format: json batch_size: 64

使用PyYAML加载：

import yaml with open("config.yaml", "r", encoding="utf-8") as f: config = yaml.safe_load(f)

5. 总结

通过对推理.py脚本的系统性改造，我们可以显著提升 MGeo 在真实业务场景中的适应能力。本文介绍了从输入输出格式解耦、多数据源支持、错误容错机制到批量性能优化的完整实践路径。

核心要点总结如下：

1. 解耦输入输出逻辑：通过参数化设计实现格式自由切换。
2. 增强兼容性：支持 JSON/CSV 等常见格式，并能自动识别字段别名。
3. 丰富输出内容：加入状态码、耗时、版本等元信息，便于监控与追溯。
4. 提升鲁棒性：引入异常处理与日志记录，保障长时间运行稳定。
5. 面向生产优化：支持批量推理与配置驱动，降低运维成本。

这些修改不仅适用于地址相似度任务，也可作为通用 NLP 推理脚本的工程范本，帮助开发者快速构建高可用的服务接口。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。