from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype=torch.float16,
    device_map="auto"
)

model = PeftModel.from_pretrained(
    base_model, 
    "nvidia/llama-3.1-nemoguard-8b-content-safety"
)
model = model.merge_and_unload()

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

repo_id = "oneonlee/llama-3.1-nemoguard-8b-content-safety-merged" 

model.push_to_hub(
    repo_id,
    use_auth_token=True,
    commit_message="Merged Llama-3.1-8B with NemoGuard content safety adapter",
    private=False, 
)
tokenizer.push_to_hub(repo_id, use_auth_token=True)


from huggingface_hub import HfApi

api = HfApi()
readme_content = """
# Llama-3.1-NemoGuard-8B Content Safety Merged

This model is a merged version of:
- Base: meta-llama/Llama-3.1-8B-Instruct
- Adapter: nvidia/llama-3.1-nemoguard-8b-content-safety

## Usage
```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("oneonlee/llama-3.1-nemoguard-8b-content-safety-merged")
tokenizer = AutoTokenizer.from_pretrained("oneonlee/llama-3.1-nemoguard-8b-content-safety-merged")
```
"""

api.upload_file(
    path_or_fileobj=readme_content.encode(),
    path_in_repo="README.md",
    repo_id=repo_id,
    repo_type="model",
)