Browse Source

[to #42322933] Fix bug for text generation task model

Fixed the bug for generate method in TaskModelForTextGeneration, which was unavailable due to the upgrade of the transformers library to version 4.24.0

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/10791805
master^2
hemu.zp wenmeng.zwm 3 years ago
parent
commit
8184c86c5f
1 changed files with 2 additions and 3 deletions
  1. +2
    -3
      modelscope/models/nlp/task_models/text_generation.py

+ 2
- 3
modelscope/models/nlp/task_models/text_generation.py View File

@@ -2,7 +2,7 @@
from typing import Any, Dict

import numpy as np
from transformers.modeling_utils import GenerationMixin
from transformers.modeling_utils import PreTrainedModel

from modelscope.metainfo import TaskModels
from modelscope.models.builder import MODELS
@@ -17,8 +17,7 @@ __all__ = ['TaskModelForTextGeneration']

@MODELS.register_module(
Tasks.text_generation, module_name=TaskModels.text_generation)
class TaskModelForTextGeneration(SingleBackboneTaskModelBase, GenerationMixin):
main_input_name = 'input_ids'
class TaskModelForTextGeneration(SingleBackboneTaskModelBase, PreTrainedModel):

def __init__(self, model_dir: str, *args, **kwargs):
"""initialize the text generation model from the `model_dir` path.


Loading…
Cancel
Save