Views
No views yet
1import torch
2
3from transformers import AutoTokenizer, TextStreamer, AutoModelForCausalLM
4
5model_path = "Crystalcareai/GemMoE-Medium-v0.5"
6
7# Load model
8
9model = AutoModelForCausalLM.from_pretrained(
10 model_path,
11 device_map="auto",
12 low_cpu_mem_usage=True,
13 torch_dtype=torch.float16,
14 attn_implementation="flash_attention_2",
15 trust_remote_code=True,
16)
17
18tokenizer = AutoTokenizer.from_pretrained(model_path)
19
20streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
21
22# Convert prompt to tokens
23
24prompt_template = "[INST] {prompt} [/INST]"
25prompt = "You're standing on the surface of the Earth. "\
26 "You walk one mile south, one mile west and one mile north. "\
27 "You end up exactly where you started. Where are you?"
28
29tokens = tokenizer(
30 prompt_template.format(prompt=prompt),
31 return_tensors='pt'
32).input_ids.cuda()
33
34# Generate output
35
36generation_output = model.generate(
37 tokens,
38 streamer=streamer,
39 max_new_tokens=512
40)