MLflow Sentence Transformers Flavor 完整指南模型记录、加载与部署语义嵌入服务【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflowMLflow 的sentence_transformersflavor 为 Sentence Transformers 库提供了原生集成用于从文本生成语义嵌入semantic embeddings。本文以mlflow/sentence_transformers/__init__.py实现为骨架系统讲解如何将句子嵌入模型保存为 MLflow Model、纳入实验跟踪与模型注册、以原生接口或 PyFunc 接口加载推理并支持以 OpenAI 兼容的llm/v1/embeddings任务形式部署为标准化嵌入服务同时覆盖签名、依赖环境、Spark UDF 等实战细节。概览什么是 MLflow Sentence Transformers FlavorSentence Transformers 是目前构建语义搜索、语义相似度、文本聚类与文本分类等应用的主流工具其核心 APISentenceTransformer.encode()将文本映射为固定维度的稠密向量。MLflow 的sentence_transformersflavor实现文件为 mlflow/sentence_transformers/init.py在此基础上提供了一层工程化封装核心能力包括模型记录与版本化通过save_model/log_model将训练好的SentenceTransformer模型连同完整元数据、输入输出签名、依赖环境一并序列化供 MLflow Tracking 与 Model Registry 管理。统一部署接口以 PyFuncPython Function方式加载模型获得统一的predict()推理接口可直接部署为 REST 服务或用于 Spark 批处理。标准化嵌入任务支持llm/v1/embeddings推理任务将模型输出包装为 OpenAI 兼容的嵌入响应格式含object、data、usage字段。依赖与环境自动管理自动生成requirements.txt、conda.yaml、python_env.yaml保障模型可复现运行。安装方式为pip install mlflow[sentence-transformers]安装后即可通过mlflow.sentence_transformers模块使用本 flavor 的全部 API官方使用指南见 docs/docs/classic-ml/deep-learning/sentence-transformers/index.mdx。模型保存save_model 参数与底层行为sentence_transformers.save_model()负责将一个训练好的 SentenceTransformer 模型持久化到本地路径。其完整参数定义源码见 mlflow/sentence_transformers/init.py如下参数类型说明modelSentenceTransformer待保存的训练好的 sentence-transformers 模型pathstr序列化模型的本地保存路径taskstr \| NoneMLflow 推理任务类型当前仅支持llm/v1/embeddingsinference_configdict \| None推理时应用的参数字典仅在以 PyFunc 或 Spark 方式加载时生效不影响load_model()返回的原生模型code_pathslist[str] \| None需随模型一并保存的自定义代码文件路径mlflow_modelModel \| None指定一个已存在的 MLflow Model 对象缺省时自动创建signatureModelSignature \| None描述模型输入输出的签名缺省但有input_example时自动推断两者都缺省时使用默认签名传False可阻止任何签名被写入input_exampleModelInputExample \| None输入示例用于校验与自动推断签名pip_requirementslist[str] \| str \| Nonepip 依赖列表或指向依赖文件的路径extra_pip_requirementslist[str] \| str \| None在默认依赖基础上额外追加的 pip 依赖conda_envdict \| None自定义 conda 环境缺省时自动生成metadatadict \| None附加到模型上的自定义元数据extra_fileslist[str] \| None需随模型额外保存的文件kwargsdict透传给model.save()的额外参数保存流程的源码级细节从 save_model 的实现可以看到几个关键行为模型数据路径固定模型权重被保存到保存路径下的model.sentence_transformer子目录常量SENTENCE_TRANSFORMERS_DATA_PATH model.sentence_transformer定义于模块顶部第 62 行。签名优先级当显式传入task时签名会被强制替换为 LLM 嵌入任务的输入输出 Schema否则优先基于input_example自动推断再否则回退到_get_default_signature()该默认签名声明输入为单个string列、输出为float64的[-1]维度张量第 503-511 行对应测试 test_default_signature_assignment 有精确断言。双 flavor 注册pyfunc.add_to_model()注册python_functionflavorloader 为mlflow.sentence_transformers同时mlflow_model.add_flavor(FLAVOR_NAME, ...)注册sentence_transformersflavor 并记录sentence_transformers_version、底层 Transformers 模型的source_model_name与pipeline_model_type等元数据第 212-230 行。测试 test_logged_data_structure 验证了MLmodel文件中这些字段的真实形态例如all-MiniLM-L6-v2会记录source_model_name: sentence-transformers/all-MiniLM-L6-v2、pipeline_model_type: BertModel。环境文件生成若未提供conda_env会先通过infer_pip_requirements推断依赖并与get_default_pip_requirements()的默认依赖sentence-transformers、transformers、torch求并集后写出conda.yaml、requirements.txt必要时含constraints.txt与python_env.yaml四个环境文件第 238-259 行。inference_config 落地如果传入了inference_config会以 JSON 形式写入保存目录下的inference_config文件供 PyFunc 加载时读取第 235-236 行。模型记录log_model 与实验跟踪log_model()是save_model()的 Tracking 集成版本将模型作为当前 run 的 artifact 记录并可选地注册到 Model Registry。其参数在save_model基础上增加了registered_model_name注册的模型名、await_registration_for等待注册完成的秒数默认 5 分钟传 0 或None表示不等待以及prompts、name、params、tags、model_type、step、model_id等通用记录参数第 302-429 行。官方文档给出了一段可直接运行的示例log_model docstringfrom sentence_transformers import SentenceTransformer import mlflow model SentenceTransformer(all-MiniLM-L6-v2) data MLflow is awesome! signature mlflow.models.infer_signature( model_inputdata, model_outputmodel.encode(data), ) with mlflow.start_run(): mlflow.sentence_transformers.log_model( modelmodel, namesbert_model, signaturesignature, input_exampledata, )在 docs/docs/classic-ml/deep-learning/sentence-transformers/index.mdx 中还给出了结合实验跟踪的完整用法在一个 run 内log_params记录模型名、嵌入维度、语料规模等超参数log_artifact保存语料文件最后log_model记录模型形成训练—跟踪—产物归档一体化的实验流程。log_model的调用链最终进入mlflow.models.Model.log()将save_model之外的所有参数一并透传第 406-429 行。测试 test_log_model_calls_register_model 证实传入registered_model_name时会自动调用注册接口而不传则不会注册。模型加载原生接口与 PyFunc 接口load_model加载原生 SentenceTransformerimport mlflow model_info ... # log_model 返回的 ModelInfo包含 model_uri loaded mlflow.sentence_transformers.load_model(model_info.model_uri) embeddings loaded.encode([Hello world, MLflow is great])load_model()第 462-500 行支持多种 URI 形式本地路径/Users/me/path/to/local/model、相对路径、s3://、runs:/run_id/path、mlflow-artifacts:/path/to/model等。加载过程会先调用_download_artifact_from_uri下载/解析 artifact再定位到model.sentence_transformer子目录用SentenceTransformer(local_model_dir)构造原生模型返回的是可直接使用.encode()的sentence_transformers实例——注意inference_config对这种方式加载的模型不生效见参数说明。一个值得注意的细节是_get_load_kwargs()第 432-443 行当 sentence-transformers 版本 ≥ 2.3.0 时加载会自动设置trust_remote_codeTrue。源码注释解释了原因MLflow 保存模型时已将整个仓库文件随 artifact 一起保存因此不存在运行远端未信任代码的风险除非 artifact 本身被恶意修改——那已经是更广义的安全范畴。测试 test_model_save_and_load_with_custom_code 验证了Alibaba-NLP/gte-base-en-v1.5这类依赖远程代码的模型也能正确保存与加载。PyFunc统一的部署接口pyfunc_model mlflow.pyfunc.load_model(model_info.model_uri) result pyfunc_model.predict([Hello world, MLflow is great])_load_pyfunc()第 446-459 行返回_SentenceTransformerModelWrapper包装对象其predict()方法第 525-565 行对输入做了大量兼容处理输入形态归一化单条字符串、字符串列表、pd.Series、pd.DataFrame、numpy 数组均可输入统一转换为列表后交给model.encode()sentence-transformers v5.4 起encode拒绝DataFrame/Series直接输入因此必须先转换对应上游 huggingface/sentence-transformers#3554。推理参数透传params会作为encode()的 kwargs 透传如batch_size若参数非法则抛出MlflowExceptionReceived invalid parameter value forparamsargument。但注意只有当模型签名定义了 params schema 时params才会在推理期生效否则会被忽略并打印警告详见测试 test_model_pyfunc_predict_with_params。任务感知输出若模型以llm/v1/embeddings任务保存且输入为{input: ...}形式的字典则输出会被包装为 OpenAI 兼容格式。测试 test_model_pyfunc_save_load 验证了单句与多句输入在 PyFunc 下的输出形状分别为(1, embed_dim)与(n, embed_dim)。llm/v1/embeddings 任务将模型部署为标准化嵌入服务这是本 flavor 面向现代 AI 应用Agent、RAG 等的关键能力。通过taskllm/v1/embeddings保存模型后模型的输入输出被固定为 LLM 嵌入 API 的标准契约可直接对接 OpenAI 兼容的客户端生态。任务校验与元数据mlflow.sentence_transformers.save_model(model, path, taskllm/v1/embeddings)当传入task时_verify_task_and_update_metadata()第 96-112 行会校验task只能是llm/v1/embeddings其他值直接抛出MlflowExceptionTask type could only be llm/v1/embeddings若metadata中已存在task且与传入值不一致也会报错Task type is inconsistent with the task value from metadata。校验通过后task被写入模型 metadata同时model_config中记录{task: llm/v1/embeddings}第 206-208 行PyFunc 加载时从model_config读取该任务信息第 457-459 行。相关校验逻辑均有测试覆盖test_verify_task_and_update_metadata。标准输入输出 Schema指定任务后签名会被强制替换为 mlflow/types/llm.py 中定义的 Schema输入input列string类型单条文本或文本列表输出包含objectstring、data数组元素为含index、object、embedding属性的对象、usage含prompt_tokens、total_tokens三个字段。推理时的输出包装当 PyFunc 收到{input: ...}形式输入且模型带llm/v1/embeddings任务时包装器会调用postprocess_output_for_llm_v1_embedding_task()定义于 mlflow/transformers/llm_inference_utils.py输出形如{ object: list, data: [ {object: embedding, index: 0, embedding: [0.4713, 0.4669, ...]} ], usage: {prompt_tokens: 8, total_tokens: 8} }其中prompt_tokens由 tokenizer 对输入逐条统计得到。测试 test_model_pyfunc_with_dict_input 验证了单条与多条输入下输出结构、嵌入维度与 token 统计的正确性。模型签名与输入示例签名signature是模型部署时的契约决定请求校验与推理参数的处理方式。本 flavor 遵循以下优先级规则save_model 逻辑显式传入task→ 使用 LLM 嵌入任务 Schema传入input_example→ 基于示例自动推断签名两者都缺省 → 使用默认签名string→float64[-1]张量显式传signatureFalse→ 不写入任何签名。推荐在生产部署时显式构造签名from mlflow.models import infer_signature sample_texts [ MLflow makes ML development easier, Sentence transformers create embeddings, ] sample_embeddings model.encode(sample_texts) signature infer_signature(sample_texts, sample_embeddings) with mlflow.start_run(): mlflow.sentence_transformers.log_model( modelmodel, namemodel, signaturesignature, input_examplesample_texts, )测试 test_signature_and_examples_are_saved_correctly 系统验证了四种组合无示例无签名/有示例/显式签名/两者皆有下保存的签名与输入示例与预期一致且input_example为 DataFrame 时以dataframe类型保存、为字符串列表时以json_object类型保存。依赖环境管理模型的可复现性是工程化的关键。本 flavor 的默认依赖为sentence-transformers、transformers、torch三个基础包由get_default_pip_requirements()第 73-84 行以固定版本pinned形式返回get_default_conda_env()第 87-93 行则在此基础上生成默认 conda 环境。测试 test_dependency_mapping 确认了默认依赖集合。实际使用时可通过三种方式定制环境# 方式一pip_requirements 覆盖全部依赖列表或文件路径 mlflow.sentence_transformers.log_model( model, namemodel, pip_requirements[sentence-transformers2.7.0, torch2.3.0] ) # 方式二extra_pip_requirements 在默认依赖基础上追加 mlflow.sentence_transformers.log_model( model, namemodel, extra_pip_requirements[my-custom-package] ) # 方式三conda_env 指定完整 conda 环境文件 mlflow.sentence_transformers.log_model(model, namemodel, conda_envpath/to/conda.yaml)测试 test_log_with_pip_requirements 与 test_log_with_extra_pip_requirements 验证了三种方式及-r引用文件、-c约束文件语法的行为差异pip_requirements是严格替换默认依赖而extra_pip_requirements是追加。此外若未显式指定环境MLflow 会先推断实际依赖再与默认依赖求并集确保模型在其真实所需环境中可运行第 240-242 行。服务化部署与 Spark 批处理部署为 REST 服务PyFunc 模型可直接启动为 REST 服务。测试 test_pyfunc_serve_and_score 展示了完整流程先用input_example记录模型再用mlflow.pyfunc的评分服务器以 JSON 载荷{inputs: [...]}请求服务并断言本地predict()与线上服务结果完全一致——这意味着你可以在本地 PyFunc 与线上服务之间获得完全确定的行为。Spark UDF借助mlflow.pyfunc.spark_udf()SentenceTransformer 模型可直接注册为 Spark UDF 对 DataFrame 批量生成嵌入测试 test_spark_udffrom pyspark.sql.types import ArrayType, DoubleType import mlflow with mlflow.start_run(): signature infer_signature(sentences, model.encode(sentences), {batch_size: 16}) model_info mlflow.sentence_transformers.log_model(model, namemy_model, signaturesignature) loaded_model mlflow.pyfunc.spark_udf( spark, model_info.model_uri, result_typeArrayType(DoubleType()), params{batch_size: 16} ) df spark.createDataFrame([(hello MLflow,), (bye world,)], [text]) df df.withColumn(embedding, loaded_model(text))该示例中params传入batch_size用于控制批量编码的批大小result_type声明为ArrayType(DoubleType())以接收每行一个嵌入向量。典型应用场景将上述能力组合起来即可覆盖 Sentence Transformers 在 MLflow 生态中的典型工作流语义搜索在 run 内记录语料嵌入、log_artifact保存语料 CSV、log_model记录检索模型官方指南 Semantic Search 示例。语义相似度 / 文本聚类 / 文本分类模型输出固定维度的嵌入向量配合自定义下游逻辑使用。微调实验管理在 run 内用log_params记录基础模型、epoch、batch size、学习率等超参微调后log_model记录模型Fine-tuning 示例。RAG / Agent 嵌入服务以taskllm/v1/embeddings部署直接对接 OpenAI 兼容的嵌入 API 客户端。大规模离线批处理通过 Spark UDF 在分布式环境中批量计算嵌入。仓库还提供了多个可运行的教程 Notebook 作为参考quickstart、semantic-search、semantic-similarity、paraphrase-mining可对照 tests/sentence_transformers/test_sentence_transformers_model_export.py 中的测试用例深入理解各 API 的行为细节。【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考