baseten/example-Meta-Llama-3-8B-InstructForSequenceClassification

#!/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()