Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3def generate_prompt(instruction, input=""):
4 instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
5 input = input.strip().replace('\r\n','\n').replace('\n\n','\n')
6 if input:
7 return f"""Instruction: {instruction}
8Input: {input}
9Response:"""
10 else:
11 return f"""User: {instruction}
12
13Assistant:"""
14#model = AutoModelForCausalLM.from_pretrained("TimeMobius/Mobius-RWKV-Chat-12B-128k-v4-HF", trust_remote_code=True, torch_dtype=torch.bfloat16).to(0)
15model = AutoModelForCausalLM.from_pretrained("TimeMobius/Mobius-RWKV-Chat-12B-128k-v4-HF", trust_remote_code=True, torch_dtype=torch.float16).to(0)
16tokenizer = AutoTokenizer.from_pretrained("TimeMobius/Mobius-RWKV-Chat-12B-128k-v4-HF", trust_remote_code=True)
17text = "Write a beginning of sci-fi novel"
18prompt = generate_prompt(text)
19inputs = tokenizer(prompt, return_tensors="pt").to(0)
20output = model.generate(inputs["input_ids"], max_new_tokens=128, do_sample=True, temperature=1.0, top_p=0.3, top_k=0, )
21print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))