| ```python | |
| #!/usr/bin/env python | |
| import torch | |
| from transformers import ( | |
| AutoConfig, | |
| AutoTokenizer, | |
| AutoModelForCausalLM, | |
| LlamaForSequenceClassification, | |
| ) | |
| # install torch, transformers, accelerate | |
| def main(): | |
| # Define the input and output repository names. | |
| input_model_id = "meta-llama/Meta-Llama-3-8B-Instruct" | |
| split_2 = input_model_id.split("/")[1] | |
| output_model_id = f"baseten/example-{split_2}ForSequenceClassification" | |
| # Load the original configuration. | |
| # (If needed, add trust_remote_code=True for custom implementations.) | |
| config = AutoConfig.from_pretrained(input_model_id) | |
| # Update the config for a sequence classification task with 10 labels. | |
| num_labels = 30 | |
| config.num_labels = num_labels | |
| config.id2label = {i: f"token activation {i}" for i in range(num_labels)} | |
| config.label2id = {f"token activation {i}": i for i in range(num_labels)} | |
| # Download the tokenizer from the original model. | |
| tokenizer = AutoTokenizer.from_pretrained(input_model_id) | |
| # Load the original causal LM model. | |
| lm_model = AutoModelForCausalLM.from_pretrained(input_model_id, config=config, device_map="auto", low_cpu_mem_usage=True) | |
| config.architectures = ["LlamaForSequenceClassification"] | |
| del lm_model.model | |
| print("loaded lm model") | |
| # Initialize the sequence classification model. | |
| # NOTE: We are using the built-in LlamaForSequenceClassification, | |
| # which uses a `.score` attribute as the output head. | |
| seq_cls_model = LlamaForSequenceClassification.from_pretrained(input_model_id, config=config, device_map="auto", low_cpu_mem_usage=True) | |
| # --- Initialize the Classification Head --- | |
| # Here we re-use the first 10 rows from the original LM head | |
| # (i.e. rows 0 to 9) to initialize the new classification head. | |
| with torch.no_grad(): | |
| # lm_model.lm_head.weight has shape [vocab_size, hidden_size] | |
| # We take the first 10 rows to form a [10, hidden_size] weight matrix. | |
| seq_cls_model.score.weight.copy_(lm_model.lm_head.weight.data[:num_labels, :]) | |
| if lm_model.lm_head.bias is not None: | |
| seq_cls_model.score.bias.copy_(lm_model.lm_head.bias.data[:num_labels]) | |
| # Optionally, save the new model locally. | |
| # save_directory = f"./{output_model_id.replace('/','_')}" | |
| # seq_cls_model.save_pretrained(save_directory) | |
| # tokenizer.save_pretrained(save_directory) | |
| # Push the new model and tokenizer to the Hub. | |
| # (Ensure you are authenticated with Hugging Face Hub via `huggingface-cli login`.) | |
| tokenizer.push_to_hub(output_model_id) | |
| seq_cls_model.push_to_hub(output_model_id) | |
| print(f"New model pushed to the Hub: {output_model_id}") | |
| if __name__ == "__main__": | |
| main() | |
| ``` |