import sys
import os
from transformers import AutoModel, AutoTokenizer
from transformers.utils import cached_file

# Load model and tokenizer from Hugging Face Hub
print("Loading model and tokenizer...")
model = AutoModel.from_pretrained("hemantn/ablang2", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("hemantn/ablang2", trust_remote_code=True)

# Find the cached model directory and import adapter
adapter_path = cached_file("hemantn/ablang2", "adapter.py")
cached_model_dir = os.path.dirname(adapter_path)
sys.path.insert(0, cached_model_dir)

# Import and create the adapter
from adapter import AbLang2PairedHuggingFaceAdapter
ablang = AbLang2PairedHuggingFaceAdapter(model=model, tokenizer=tokenizer)

# Test sequences from the notebook
test_sequences = [
    ['EVQ***SGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCAR**PGHGAAFMDVWGTGTTVTVSS',
     'DIQLTQSPLSLPVTLGQPASISCRSS*SLEASDTNIYLSWFQQRPGQSPRRLIYKI*NRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK']
]

print("Testing restore without alignment:")
result_no_align = ablang(test_sequences, mode='restore', align=False)
print(f"Result (no align): {result_no_align[0]}")

print("\nTesting restore with alignment:")
result_with_align = ablang(test_sequences, mode='restore', align=True)
print(f"Result (with align): {result_with_align[0]}")

print("\nBoth options work correctly!")