初始化项目,由ModelHub XC社区提供模型

Model: AI-ModelScope/opt-125
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-08-09 13:47:13 +08:00
commit 0e390568d5
14 changed files with 50337 additions and 0 deletions

56
ms_wrapper.py Normal file
View File

@@ -0,0 +1,56 @@
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