Views
No views yet

1pip install -r requirements.txt
2pip install transformers datasets accelerate safetensors1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
3
4
5device = "cuda" if torch.cuda.is_available() else "cpu"
6
7model_id = "DireDreadlord/Dragon-1.5-0.5B"
8model = AutoModelForCausalLM.from_pretrained(
9 model_id,
10 device_map="auto",
11 dtype="auto"
12)
13model.to(device)
14tokenizer = AutoTokenizer.from_pretrained(model_id)
15streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
16
17prompt = "Solve the leetcode problem 1: two sum using the hash map technique"
18
19input_ids = tokenizer.apply_chat_template(
20 [{"role": "user", "content": prompt}],
21 add_generation_prompt=True,
22 return_tensors="pt",
23 tokenize=True,
24)["input_ids"].to(device)
25
26output = model.generate(
27 input_ids,
28 do_sample=True,
29 temperature=0.4,
30 top_k=50,
31 repetition_penalty=1.05,
32 max_new_tokens=2048,
33 streamer=streamer,
34)max_new_tokens=2048