# 轻量级多任务模型架构
class LightweightExperimentModel:
    def __init__(self):
        # 基础backbone（支持多种预训练模型）
        self.backbone = DistilBertModel.from_pretrained(
            'distilbert-base-chinese'
        )
        
        # 特征压缩层（768 -> 256维）
        self.feature_compressor = nn.Sequential(
            nn.Linear(768, 256),
            nn.ReLU(),
            nn.Dropout(0.1)
        )
        
        # 多任务头设计
        self.op_classifier = nn.Linear(256, 3)     # 操作分类
        self.entity_detector = nn.Linear(256, 5)   # 实体识别
        self.param_parser = nn.Linear(256, 2)      # 参数解析
        
        # CRF层用于序列标注
        self.crf = CRF(num_tags=5, batch_first=True)
        
    def forward(self, inputs):
        # 特征提取与压缩
        features = self.backbone(inputs).last_hidden_state
        compressed = self.feature_compressor(features)
        
        # 多任务预测
        return {
            'operation': self.op_classifier(compressed),
            'entities': self.entity_detector(compressed),
            'parameters': self.param_parser(compressed)
        }