import torch, json
from transformers import BertTokenizer, BertForSequenceClassification
from huggingface_hub import hf_hub_download

REPO = "Aniruddhhhh/jungian-dream-archetype-group_2"
tokenizer = BertTokenizer.from_pretrained(REPO)
model     = BertForSequenceClassification.from_pretrained(REPO)
model.eval()

arch_path  = hf_hub_download(REPO, "archetypes.json")
with open(arch_path) as f:
    archetypes = json.load(f)

dream = "I was chased through a dark forest by a shadowy figure."
enc   = tokenizer(dream, return_tensors="pt",
                  truncation=True, max_length=256)
with torch.no_grad():
    probs = torch.sigmoid(model(**enc).logits)[0]
for arch, prob in zip(archetypes, probs):
    print(f"{arch}: {prob:.3f}")