| ```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() |
| |
| ``` |