import argparse
import json
from pathlib import Path

from huggingface_hub import snapshot_download
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


def load_adapter(subfolder, revision=None):
    repo_id = 'SJTUFreshman/collaborativelearning'
    tokenizer_folder = str(Path(subfolder).parent).replace('\\', '/') if subfolder.endswith('/policy_adapter') else subfolder
    snapshot = snapshot_download(repo_id, revision=revision, allow_patterns=[subfolder+'/*', tokenizer_folder+'/tokenizer*', tokenizer_folder+'/chat_template.jinja', tokenizer_folder+'/special_tokens_map.json'])
    adapter_dir = Path(snapshot)/subfolder
    config = json.loads((adapter_dir/'adapter_config.json').read_text())
    tokenizer = AutoTokenizer.from_pretrained(Path(snapshot)/tokenizer_folder)
    base = AutoModelForCausalLM.from_pretrained(config['base_model_name_or_path'], torch_dtype='auto', device_map='auto')
    model = PeftModel.from_pretrained(base, adapter_dir)
    return model, tokenizer


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('subfolder')
    parser.add_argument('--revision')
    args = parser.parse_args()
    model, tokenizer = load_adapter(args.subfolder, args.revision)
    print('Loaded', args.subfolder)