144 lines
5.2 KiB
Python
144 lines
5.2 KiB
Python
import os
|
||
import torch
|
||
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.outputs import OutputKeys
|
||
from modelscope.pipelines.nlp.text_generation_pipeline import TextGenerationPipeline
|
||
from modelscope.models.base import Model, TorchModel
|
||
from modelscope.utils.logger import get_logger
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
|
||
from transformers.generation.utils import GenerationConfig
|
||
|
||
import torch
|
||
|
||
from modelscope.models.base import TorchModel
|
||
from modelscope.preprocessors.base import Preprocessor
|
||
from modelscope.pipelines.base import Model, Pipeline
|
||
from modelscope.utils.config import Config
|
||
from modelscope.pipelines.builder import PIPELINES
|
||
from modelscope.preprocessors.builder import PREPROCESSORS
|
||
from modelscope.models.builder import MODELS
|
||
# m_path="你的路径/cvx-coder"
|
||
# model = AutoModelForCausalLM.from_pretrained(
|
||
# m_path,
|
||
# device_map="cuda",
|
||
# torch_dtype="auto",
|
||
# trust_remote_code=True,
|
||
# )
|
||
# tokenizer = AutoTokenizer.from_pretrained(m_path)
|
||
# pipe = pipeline(
|
||
# "text-generation",
|
||
# model=model,
|
||
# tokenizer=tokenizer,
|
||
# )
|
||
# generation_args = {
|
||
# "max_new_tokens": 2000,
|
||
# "return_full_text": False,
|
||
# "temperature": 0,
|
||
# "do_sample": False,
|
||
# }
|
||
# content='''my problem is not convex, can i use cvx? if not, what should i do, be specific.'''
|
||
# messages = [
|
||
# {"role": "user", "content": content},
|
||
# ]
|
||
# output = pipe(messages, **generation_args)
|
||
# print(output[0]['generated_text'])
|
||
|
||
@PIPELINES.register_module('text-generation', module_name='cvx-coder-pipe')
|
||
class Baichuan7BTextGenerationPipeline(TextGenerationPipeline):
|
||
def __init__(
|
||
self,
|
||
model: Union[Model, str],
|
||
*args,
|
||
**kwargs):
|
||
self.model = Baichuan7BTextGeneration(model) if isinstance(model, str) else model
|
||
super().__init__(model=model, **kwargs)
|
||
|
||
def preprocess(self, inputs, **preprocess_params) -> Dict[str, Any]:
|
||
return inputs
|
||
|
||
def _sanitize_parameters(self, **pipeline_parameters):
|
||
return {},pipeline_parameters,{}
|
||
|
||
# define the forward pass
|
||
def forward(self, inputs: str, **forward_params) -> Dict[str, Any]:
|
||
output = {}
|
||
content=inputs
|
||
messages = [
|
||
{"role": "user", "content": content},
|
||
]
|
||
outputs = self.model.pipeline(messages, **self.model.generation_args)
|
||
|
||
output['text'] = outputs[0]['generated_text']
|
||
return output
|
||
|
||
# format the outputs from pipeline
|
||
def postprocess(self, input, **kwargs) -> Dict[str, Any]:
|
||
return input
|
||
|
||
|
||
@MODELS.register_module('text-generation', module_name='cvx-coder')
|
||
class Baichuan7BTextGeneration(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)
|
||
self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", torch_dtype="auto", trust_remote_code=True)
|
||
# self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto",trust_remote_code=True)
|
||
self.model.generation_config = GenerationConfig.from_pretrained(model_dir)
|
||
self.model = self.model.eval()
|
||
from transformers import pipeline
|
||
self.pipeline=pipeline("text-generation",model=self.model,tokenizer=self.tokenizer,)
|
||
self.generation_args= {"max_new_tokens": 2000, "return_full_text": False,"temperature": 0, "do_sample": False,}
|
||
|
||
def forward(self, input: str, *args, **kwargs) -> Dict[str, Any]:
|
||
content = input
|
||
messages = [
|
||
{"role": "user", "content": content},
|
||
]
|
||
outputs = self.model.pipeline(messages, **self.model.generation_args)
|
||
|
||
response = outputs[0]['generated_text']
|
||
return {OutputKeys.RESPONSE:response, OutputKeys.HISTORY: ""}
|
||
|
||
def quantize(self, bits: int):
|
||
self.model = self.model.quantize(bits)
|
||
return self
|
||
|
||
def infer(self, input, **kwargs):
|
||
content = input
|
||
messages = [
|
||
{"role": "user", "content": content},
|
||
]
|
||
outputs = self.pipeline(messages, **self.model.generation_args)
|
||
|
||
response = outputs[0]['generated_text']
|
||
return response
|
||
|
||
|
||
# Tips: usr_config_path is the temporary save configuration location， after upload modelscope hub, it is the model_id
|
||
# usr_config_path = '/mnt/workspace/cvx-coder3'
|
||
usr_config_path = './'
|
||
|
||
config = Config({
|
||
"framework": 'pytorch',
|
||
"task": 'text-generation',
|
||
"model": {'type': 'cvx-coder'},
|
||
"pipeline": {"type": "cvx-coder-pipe"},
|
||
"allow_remote": True
|
||
})
|
||
config.dump('./'+ 'configuration.json')
|
||
|
||
if __name__ == "__main__":
|
||
from modelscope.models import Model
|
||
from modelscope.pipelines import pipeline
|
||
# model = Model.from_pretrained(usr_config_path)
|
||
input = "Hello, ModelScope!"
|
||
inference = pipeline('text-generation', model=usr_config_path)
|
||
output = inference(input)
|
||
print(output)
|