70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
|
# from modelscope.utils.constant import Tasks
|
|||
|
|
# from modelscope.pipelines import pipeline
|
|||
|
|
|
|||
|
|
# prompt = "输入: 中国的首都在哪里\n输出: "
|
|||
|
|
|
|||
|
|
# # task可选值为 抽取、分类。text为需要分析的文本。labels为类型列表，中文逗号分隔。
|
|||
|
|
# inputs = {'task': '抽取', 'text': '杭州欢迎你。', 'labels': '地名'}
|
|||
|
|
# # PROMPT_TEMPLATE保持不变
|
|||
|
|
# PROMPT_TEMPLATE = '输入: {text}\n{task}: {labels}\n输出: '
|
|||
|
|
# # prompt = PROMPT_TEMPLATE.format(**inputs)
|
|||
|
|
# pipeline_ins = pipeline(task=Tasks.text_generation, model='damo/nlp_seqgpt-560m', model_revision = 'v1.0.1', run_kwargs={'gen_token': '[GEN]'})
|
|||
|
|
# print(pipeline_ins(prompt))
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import json
|
|||
|
|
from modelscope.pipelines import pipeline
|
|||
|
|
from modelscope.utils.constant import Tasks
|
|||
|
|
|
|||
|
|
from modelscope.preprocessors import TextGenerationTransformersPreprocessor
|
|||
|
|
|
|||
|
|
from modelscope.utils.config import Config
|
|||
|
|
from modelscope import snapshot_download
|
|||
|
|
from transformers import AutoConfig
|
|||
|
|
|
|||
|
|
from modeling_bloom import BloomForCausalLM
|
|||
|
|
from tokenization_bloom_fast import BloomTokenizerFast
|
|||
|
|
|
|||
|
|
seed = 4321
|
|||
|
|
torch.manual_seed(seed)
|
|||
|
|
torch.cuda.manual_seed_all(seed)
|
|||
|
|
|
|||
|
|
model_dir = snapshot_download("damo/nlp_seqgpt-560m")
|
|||
|
|
|
|||
|
|
config, kwargs = AutoConfig.from_pretrained(
|
|||
|
|
model_dir,
|
|||
|
|
return_unused_kwargs=True,
|
|||
|
|
trust_remote_code=True,
|
|||
|
|
code_revision=None,
|
|||
|
|
_commit_hash=None,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
tokenizer_config_file = "./tokenizer_config.json"
|
|||
|
|
if tokenizer_config_file is not None:
|
|||
|
|
with open(tokenizer_config_file, encoding="utf-8") as tokenizer_config_handle:
|
|||
|
|
init_kwargs = json.load(tokenizer_config_handle)
|
|||
|
|
init_kwargs.pop("tokenizer_class", None)
|
|||
|
|
init_kwargs.pop("tokenizer_file", None)
|
|||
|
|
saved_init_inputs = init_kwargs.pop("init_inputs", ())
|
|||
|
|
init_inputs = saved_init_inputs
|
|||
|
|
init_kwargs["vocab_file"] = None
|
|||
|
|
init_kwargs["added_tokens_file"] = None
|
|||
|
|
init_kwargs["special_tokens_map_file"] = "./special_tokens_map.json"
|
|||
|
|
init_kwargs["tokenizer_file"] = "./tokenizer.json"
|
|||
|
|
init_kwargs["name_or_path"] = model_dir
|
|||
|
|
tokenizer = BloomTokenizerFast(*init_inputs, **init_kwargs)
|
|||
|
|
|
|||
|
|
model = BloomForCausalLM(config)
|
|||
|
|
model = model.from_pretrained(model_dir).cuda().eval()
|
|||
|
|
|
|||
|
|
prompt = "输入: 中国的首都在哪里\n输出: "
|
|||
|
|
prompt = "输入: 美国的首都在哪里\n输出: "
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
input_ids = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=1024)
|
|||
|
|
input_ids = input_ids.input_ids.cuda()
|
|||
|
|
outputs = model.generate(input_ids, num_beams=4, do_sample=False, max_new_tokens=256)
|
|||
|
|
decoded_sentences = tokenizer.batch_decode(outputs, skip_special_tokens=True)
|
|||
|
|
print(decoded_sentences[0])
|