56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
import os
|
|
from typing import Union, Dict, Any
|
|
from modelscope.pipelines.builder import PIPELINES
|
|
from modelscope.models.builder import MODELS
|
|
from modelscope.utils.constant import Tasks
|
|
from modelscope.pipelines.base import Pipeline
|
|
from modelscope.models.base import Model, TorchModel
|
|
from modelscope.utils.logger import get_logger
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from transformers import pipeline
|
|
|
|
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
|
|
|
|
@PIPELINES.register_module(Tasks.text_generation, module_name='opt-125-text-generation-pipe')
|
|
class Opt125TextGenerationPipeline(Pipeline):
|
|
def __init__(
|
|
self,
|
|
model: Union[Model, str],
|
|
*args,
|
|
**kwargs):
|
|
model = Opt125TextGeneration(model) if isinstance(model, str) else model
|
|
super().__init__(model=model, **kwargs)
|
|
|
|
def preprocess(self, inputs, **preprocess_params) -> Dict[str, Any]:
|
|
return inputs
|
|
|
|
# define the forward pass
|
|
def forward(self, inputs: Dict, **forward_params) -> Dict[str, Any]:
|
|
return self.model(inputs)
|
|
|
|
# format the outputs from pipeline
|
|
def postprocess(self, input, **kwargs) -> Dict[str, Any]:
|
|
return input
|
|
|
|
@MODELS.register_module(Tasks.text_generation, module_name='opt-125')
|
|
class Opt125TextGeneration(TorchModel):
|
|
def __init__(self, model_dir=None, *args, **kwargs):
|
|
super().__init__(model_dir, *args, **kwargs)
|
|
self.logger = get_logger()
|
|
# loading tokenizer
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir,trust_remote_code=True)
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto",trust_remote_code=True)
|
|
self.generator = pipeline('text-generation', model=model_dir)
|
|
self.model = self.model.eval()
|
|
|
|
def forward(self,input: Dict, *args, **kwargs) -> Dict[str, Any]:
|
|
output = {}
|
|
res = self.generator(input)
|
|
output['text'] = res[0]['generated_text']
|
|
return output
|
|
|
|
def quantize(self, bits: int):
|
|
self.model = self.model.quantize(bits)
|
|
return self
|
|
|