Views
No views yet
1transformers>=4.27.1
2accelerate>=0.20.1
3torch>=2.0.1
4modelscope>=1.8.3
5sentencepiece==0.1.991import os
2os.environ["CUDA_VISIBLE_DEVICES"] = "0"
3from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
4import torch
5
6
7model_id = "wisdomOcean/wisdomInterrogatory"
8revision = 'v1.0.0'
9model_dir = snapshot_download(model_id, revision)
10
11def generate_response(prompt: str) -> str:
12 inputs = tokenizer(f'</s>Human:{prompt} </s>Assistant: ', return_tensors='pt')
13 inputs = inputs.to('cuda')
14 pred = model.generate(**inputs, max_new_tokens=800,
15 repetition_penalty=1.2)
16 response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=True)
17 return response.split("Assistant: ")[1]
18
19tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
20model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto",
21 torch_dtype=torch.float16,
22 trust_remote_code=True)
23prompt = "如果喝了两斤白酒后开车,会有什么后果?"
24resp = generate_response(prompt)
25print(resp)