Views
No views yet
1import torch
2from transformers import AutoTokenizer, TextStreamer, AutoModelForCausalLM
3
4model_path = "Crystalcareai/llama-3-4x8b"
5model = AutoModelForCausalLM.from_pretrained(
6 model_path,
7 device_map="auto",
8 low_cpu_mem_usage=True,
9 torch_dtype=torch.bfloat16,
10 trust_remote_code=True,
11 attn_implementation="flash_attention_2",
12)
13
14tokenizer = AutoTokenizer.from_pretrained(model_path)
15streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
16
17# Modify the prompt to match the Alpaca instruction template
18prompt = """
19Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
20
21### Instruction:
22Sam is faster than Joe. Joe is faster than Jane. Is Sam faster than Jane? Explain your reasoning step by step.
23
24### Input:
25
26### Response:
27"""
28
29tokens = tokenizer(
30 prompt,
31 return_tensors='pt'
32).input_ids.cuda()
33
34generation_output = model.generate(
35 tokens,
36 streamer=streamer,
37 max_new_tokens=512,
38)