Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2from transformers.generation import GenerationConfig
3
4# Note: The default behavior now has injection attack prevention off.
5tokenizer = AutoTokenizer.from_pretrained("QiYuan-tech/LLM-Detector-Small-zh", trust_remote_code=True)
6
7# use bf16
8# model = AutoModelForCausalLM.from_pretrained("QiYuan-tech/LLM-Detector-Small-zh", device_map="auto", trust_remote_code=True, bf16=True).eval()
9
10# use fp16
11# model = AutoModelForCausalLM.from_pretrained("QiYuan-tech/LLM-Detector-Small-zh", device_map="auto", trust_remote_code=True, fp16=True).eval()
12
13# use cpu only
14# model = AutoModelForCausalLM.from_pretrained("QiYuan-tech/LLM-Detector-Small-zh", device_map="cpu", trust_remote_code=True).eval()
15
16# use auto mode, automatically select precision based on the device.
17model = AutoModelForCausalLM.from_pretrained("QiYuan-tech/LLM-Detector-Small-zh", device_map="auto", trust_remote_code=True).eval()
18#model = AutoModelForCausalLM.from_pretrained("QiYuan-tech/LLM-Detector-Small-zh", device_map="auto", trust_remote_code=True).cuda()
19
20# Specify hyperparameters for generation. But if you use transformers>=4.32.0, there is no need to do this.
21# model.generation_config = GenerationConfig.from_pretrained("./Qwen-1_8B-Chat", trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参
22
23response, history = model.chat(tokenizer, "你好", history=None)
24print(response)