Spaces:
Sleeping
Sleeping
File size: 2,471 Bytes
07cb7d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """
模型层契约定义。
这个包集中放模型产物、生成上下文和模型构建器协议。
Pipeline、服务和具体模型实现都从这里引用模型相关契约。
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Protocol
import keras
@dataclass
class GenerationContext:
end_of_text: int
max_length: int
sample_fn: Callable
@dataclass
class GenerationResult:
token_ids: list[int]
stop_reason: str
GenerateFn = Callable[[GenerationContext, list[int]], GenerationResult]
@dataclass
class ModelArtifact:
model: keras.Model
@dataclass
class TextGenerationModel:
model: keras.Model
generate: GenerateFn
class SupervisedModelBuilder(Protocol):
"""有监督任务模型构建器需要提供训练、推理和加载能力。"""
def build_training_artifact(self) -> ModelArtifact:
"""构建训练产物。"""
...
def compile_training_model(self, model: keras.Model) -> None:
"""编译训练模型。"""
...
def convert_to_inference_artifact(
self,
training_artifact: ModelArtifact
) -> ModelArtifact:
"""把训练产物转换成推理产物。"""
return training_artifact
def load_inference_artifact(self, model_path: Path) -> ModelArtifact:
"""从完整模型文件加载推理产物。"""
model = keras.models.load_model(str(model_path))
return ModelArtifact(model=model)
class TextGenerationModelBuilder(Protocol):
"""文本生成模型构建器协议"""
def build_training_artifact(
self,
vocab_size: int,
sequence_length: int
) -> TextGenerationModel:
"""构建文本训练产物"""
...
def compile_training_model(self, model: keras.Model) -> None:
"""编译文本训练模型"""
...
def convert_to_inference_artifact(
self,
training_artifact: TextGenerationModel
) -> TextGenerationModel:
"""把训练产物转换成文本推理产物"""
return training_artifact
def load_inference_artifact(
self,
model_path: Path
) -> TextGenerationModel:
"""从完整模型文件加载文本推理产物"""
...
__all__ = [
"GenerateFn",
"GenerationContext",
"GenerationResult",
"ModelArtifact",
"SupervisedModelBuilder",
"TextGenerationModel",
"TextGenerationModelBuilder"
]
|