"""
Run this in your Colab notebook to push trained adapters to Hugging Face Hub.
Paste this code into a new cell at the END of your training notebook.
"""

from huggingface_hub import HfApi, notebook_login
from peft import PeftModel
from transformers import (
    XLMRobertaForSequenceClassification,
    XLMRobertaForTokenClassification,
)
import os, glob

# Login (run this first if not already logged in)
# notebook_login()  # Uncomment if needed

print("=" * 60)
print("PUSHING TRAINED ADAPTERS TO HUB")
print("=" * 60)

# ========== SENTIMENT ==========
print("\n[1/2] Sentiment Adapter")
print("-" * 40)

# Find the best checkpoint (last one with early stopping)
sent_checkpoints = sorted(glob.glob("./sent_out/checkpoint-*"))
if not sent_checkpoints:
    print("ERROR: No sentiment checkpoints found in ./sent_out/")
    print("Check: os.listdir('./sent_out')")
else:
    best_sent = sent_checkpoints[-1]  # Last checkpoint = best with early stopping
    print(f"Found checkpoint: {best_sent}")
    
    # Load base + adapter
    sent_base = XLMRobertaForSequenceClassification.from_pretrained(
        "xlm-roberta-base", num_labels=3
    )
    sent_model = PeftModel.from_pretrained(sent_base, best_sent)
    
    # Push to Hub
    print("Pushing to BrightBorn/culturaladapt-sentiment...")
    sent_model.push_to_hub("BrightBorn/culturaladapt-sentiment")
    print("✅ Sentiment adapter pushed!")

# ========== NER ==========
print("\n[2/2] NER Adapter")
print("-" * 40)

ner_checkpoints = sorted(glob.glob("./ner_out/checkpoint-*"))
if not ner_checkpoints:
    print("ERROR: No NER checkpoints found in ./ner_out/")
    print("Check: os.listdir('./ner_out')")
else:
    best_ner = ner_checkpoints[-1]
    print(f"Found checkpoint: {best_ner}")
    
    ner_base = XLMRobertaForTokenClassification.from_pretrained(
        "xlm-roberta-base", num_labels=9
    )
    ner_model = PeftModel.from_pretrained(ner_base, best_ner)
    
    print("Pushing to BrightBorn/culturaladapt-ner...")
    ner_model.push_to_hub("BrightBorn/culturaladapt-ner")
    print("✅ NER adapter pushed!")

print("\n" + "=" * 60)
print("ALL DONE!")
print("=" * 60)
print("\nVerify at:")
print("  https://huggingface.co/BrightBorn/culturaladapt-sentiment")
print("  https://huggingface.co/BrightBorn/culturaladapt-ner")
print("\nThen restart your Space:")
print("  https://huggingface.co/spaces/BrightBorn/culturaladapt-demo")