Views
No views yet
transformers.1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model_id = "your_id/your-repo-name"
5
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7model = AutoModelForCausalLM.from_pretrained(
8 model_id,
9 torch_dtype=torch.float16,
10 device_map="auto"
11)
12
13
14# chane parameter
15
16 max_seq_length = 4096
17
18## A100 cnnfig
19# 1. Configure DPO Training Arguments
20dpo_config = DPOConfig(
21 learning_rate = 1e-6,
22 per_device_train_batch_size = 4,
23 gradient_accumulation_steps = 16,
24 num_train_epochs = 1,
25 optim = "adamw_8bit",
26 weight_decay = 0.01,
27 warmup_ratio = 0.1,
28 fp16 = not is_bfloat16_supported(),
29 bf16 = is_bfloat16_supported(),
30 logging_steps = 1,
31 output_dir = "dpo_checkpoints_001_A100-2",
32 beta = 0.1,
33 max_length = 4096,
34 max_prompt_length = 2048,
35 seed = 42,
36 report_to = "none",
37)
38
39
40# Test inference
41prompt = "Your question here"
42inputs = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")
43outputs = model.generate(**inputs, max_new_tokens=512)
44print(tokenizer.decode(outputs[0]))
45