"""Use HyPE original research imports and the checkpoint hosted on the Hugging Face Hub.

Run from the original project root, so that the project's ``src`` package is importable:
  PYTHONPATH=/path/to/Hyperbolic-Prompt-Detection-and-Sanitization \
     python /path/to/kit/examples/inference_from_hub.py "Several birds flying."

Install dependencies from the official repository before running.
"""
import sys
import torch
from huggingface_hub import hf_hub_download
from transformers import CLIPTokenizer

from src.HySAC.hysac.models import HySAC
from src.HyperbolicSVDD.source.SVDD import LorentzHyperbolicOriginSVDD, project_to_lorentz

repo_id = 'HyPE-VLM/HyPE-detector'  # change if published under your personal username
prompt = sys.argv[1] if len(sys.argv) > 1 else 'Several birds that are flying together over a frozen lake.'

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LorentzHyperbolicOriginSVDD(curvature=2.3026, radius_lr=0.2, nu=0.01, center_init='origin')
weights = hf_hub_download(repo_id=repo_id, filename='best_hyperbolic_svdd_model.pth')
model.load(weights)
model.center = model.center.to(DEVICE)
hyperbolic_clip = HySAC.from_pretrained('aimagelab/hysac', device=DEVICE).to(DEVICE)
hyperbolic_clip.eval()
tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14')
input_ids = tokenizer(prompt, return_tensors='pt', padding='max_length', truncation=True, max_length=77).input_ids.to(DEVICE)
with torch.no_grad():
    embedding = hyperbolic_clip.encode_text(input_ids, project=True)
    embedding = project_to_lorentz(embedding, model.curvature)
    prediction = model.predict(embedding)
label = int(prediction.flatten()[0].item()) if torch.is_tensor(prediction) else int(prediction)
print('Benign' if label == 1 else 'Harmful')