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