Views
No views yet
1#!/usr/bin/env python
2import torch
3from transformers import (
4 AutoConfig,
5 AutoTokenizer,
6 AutoModelForCausalLM,
7 LlamaForSequenceClassification,
8)
9# install torch, transformers, accelerate
10
11def main():
12 # Define the input and output repository names.
13 input_model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
14 split_2 = input_model_id.split("/")[1]
15 output_model_id = f"baseten/example-{split_2}ForSequenceClassification"
16
17 # Load the original configuration.
18 # (If needed, add trust_remote_code=True for custom implementations.)
19 config = AutoConfig.from_pretrained(input_model_id)
20
21 # Update the config for a sequence classification task with 10 labels.
22 num_labels = 30
23 config.num_labels = num_labels
24 config.id2label = {i: f"token activation {i}" for i in range(num_labels)}
25 config.label2id = {f"token activation {i}": i for i in range(num_labels)}
26
27 # Download the tokenizer from the original model.
28 tokenizer = AutoTokenizer.from_pretrained(input_model_id)
29
30 # Load the original causal LM model.
31 lm_model = AutoModelForCausalLM.from_pretrained(input_model_id, config=config, device_map="auto", low_cpu_mem_usage=True)
32 config.architectures = ["LlamaForSequenceClassification"]
33 del lm_model.model
34 print("loaded lm model")
35 # Initialize the sequence classification model.
36 # NOTE: We are using the built-in LlamaForSequenceClassification,
37 # which uses a `.score` attribute as the output head.
38 seq_cls_model = LlamaForSequenceClassification.from_pretrained(input_model_id, config=config, device_map="auto", low_cpu_mem_usage=True)
39
40 # --- Initialize the Classification Head ---
41 # Here we re-use the first 10 rows from the original LM head
42 # (i.e. rows 0 to 9) to initialize the new classification head.
43 with torch.no_grad():
44 # lm_model.lm_head.weight has shape [vocab_size, hidden_size]
45 # We take the first 10 rows to form a [10, hidden_size] weight matrix.
46 seq_cls_model.score.weight.copy_(lm_model.lm_head.weight.data[:num_labels, :])
47 if lm_model.lm_head.bias is not None:
48 seq_cls_model.score.bias.copy_(lm_model.lm_head.bias.data[:num_labels])
49
50 # Optionally, save the new model locally.
51 # save_directory = f"./{output_model_id.replace('/','_')}"
52 # seq_cls_model.save_pretrained(save_directory)
53 # tokenizer.save_pretrained(save_directory)
54
55 # Push the new model and tokenizer to the Hub.
56 # (Ensure you are authenticated with Hugging Face Hub via `huggingface-cli login`.)
57 tokenizer.push_to_hub(output_model_id)
58 seq_cls_model.push_to_hub(output_model_id)
59
60
61 print(f"New model pushed to the Hub: {output_model_id}")
62
63if __name__ == "__main__":
64 main()
65