#!/usr/bin/env python
"""
push_adapter.py
Push the extracted LoRA adapter + projection head + processor to the Hub.
This pushes the SMALL artifacts (adapter, not the 8GB merged model).
"""
import os
from huggingface_hub import HfApi, create_repo

ADAPTER_DIR = "/workspace/models/colgranite-4.1-4b-lora-extracted"
REPO_ID     = "dineshananthi/colgranite-4.1-4b-lora-72k"   # <-- your adapter repo
PRIVATE     = False

api = HfApi()

# 1. Create the repo if it doesn't exist (no-op if it already does)
create_repo(REPO_ID, repo_type="model", private=PRIVATE, exist_ok=True)
print(f"Repo ready: {REPO_ID}")

# 2. Optional README so the page explains what this is
readme = f"""---
base_model: ibm-granite/granite-vision-4.1-4b
library_name: peft
tags:
- colbert
- late-interaction
- visual-retrieval
- colpali
---

# ColGranite 4.1-4B — LoRA Adapter

Late-interaction (ColBERT-style) visual retrieval adapter on top of
[`ibm-granite/granite-vision-4.1-4b`](https://huggingface.co/ibm-granite/granite-vision-4.1-4b).

## Contents
- `adapter_config.json`, `adapter_model.safetensors` — LoRA adapter (text-decoder projections)
- `col_linear_head.pt` — the trained 2560 to 128 projection head (REQUIRED for retrieval)
- processor / tokenizer files

## Usage
Rebuild the `ColGranite` wrapper, attach this adapter to the inner model with
`PeftModel.from_pretrained(...)`, then load `col_linear_head.pt` into the projection head.
"""
with open(os.path.join(ADAPTER_DIR, "README.md"), "w") as f:
    f.write(readme)

# 3. Upload the whole folder (adapter + head + processor + README)
print("Uploading folder...")
api.upload_folder(
    folder_path=ADAPTER_DIR,
    repo_id=REPO_ID,
    repo_type="model",
    commit_message="Add extracted LoRA adapter + projection head from checkpoint-72000",
)

print("\nDONE ->", f"https://huggingface.co/{REPO_ID}")