File size: 1,280 Bytes
0f236b2 45237dc 0f236b2 45237dc 0f236b2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
This is an MOE of Llama-3-8b with 4 experts. This does not use semantic routing, as this utilizes the deepseek-moe architecture. There is no routing, and there is no gate - all experts are active on every token.
```python
import torch
from transformers import AutoTokenizer, TextStreamer, AutoModelForCausalLM
model_path = "Crystalcareai/llama-3-4x8b"
model = AutoModelForCausalLM.from_pretrained(
model_path,
device_map="auto",
low_cpu_mem_usage=True,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained(model_path)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
# Modify the prompt to match the Alpaca instruction template
prompt = """
Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
Sam is faster than Joe. Joe is faster than Jane. Is Sam faster than Jane? Explain your reasoning step by step.
### Input:
### Response:
"""
tokens = tokenizer(
prompt,
return_tensors='pt'
).input_ids.cuda()
generation_output = model.generate(
tokens,
streamer=streamer,
max_new_tokens=512,
)
``` |