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