()
| 16 | |
| 17 | |
| 18 | def main(): |
| 19 | model_name = "Qwen/Qwen3-0.6B" |
| 20 | |
| 21 | # load the tokenizer and the model |
| 22 | tokenizer = AutoTokenizer.from_pretrained(model_name) |
| 23 | model: GenerationMixin = AutoModelForCausalLM.from_pretrained(model_name, dtype="bfloat16").to( |
| 24 | "cuda" |
| 25 | ) |
| 26 | |
| 27 | prompts = [ |
| 28 | "Give me a short introduction to large language model.", |
| 29 | "What is the capital of France?", |
| 30 | ] |
| 31 | texts = [to_chat_text(tokenizer, prompt) for prompt in prompts] |
| 32 | tokenizer.padding_side = "left" |
| 33 | model_inputs = tokenizer(texts, return_tensors="pt", padding=True).to(model.device) |
| 34 | |
| 35 | with torch.inference_mode(): |
| 36 | outputs = model.model.forward(**model_inputs) |
| 37 | |
| 38 | os.makedirs("qwen3-0.6b", exist_ok=True) |
| 39 | hidden_states = outputs.last_hidden_state[:, -1, :] # [n_hidden_states, D] |
| 40 | torch.save(hidden_states, "qwen3-0.6b/hidden_states.pt") |
| 41 | weights = model.lm_head.weight.data # [V, D] |
| 42 | torch.save(weights, "qwen3-0.6b/weights.pt") |
| 43 | |
| 44 | |
| 45 | if __name__ == "__main__": |
no test coverage detected