Spaces:
Sleeping
Sleeping
File size: 1,790 Bytes
bcfaedb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def predict_with_model(texts, bundle):
model = bundle["model"]
tokenizer = bundle["tokenizer"]
labels = bundle["labels"]
enc = tokenizer(
texts,
padding=True,
truncation=True,
max_length=200,
return_tensors="pt"
)
input_ids = enc["input_ids"].to(device)
mask = enc["attention_mask"].to(device)
with torch.no_grad():
logits, _ = model(input_ids, mask)
preds = torch.argmax(logits, dim=1)
return [labels[p] for p in preds.cpu().tolist()]
def predict_with_sdg_model(texts, bundle):
"""
Multilabel-trained model but we pick the single label with the
highest sigmoid probability for each text (argmax strategy).
Returns e.g. ["13", "7", "1", ...] -- SDG numbers as strings.
"""
model = bundle["model"]
tokenizer = bundle["tokenizer"]
labels = bundle["labels"]
enc = tokenizer(
texts,
padding=True,
truncation=True,
max_length=200,
return_tensors="pt"
)
input_ids = enc["input_ids"].to(device)
mask = enc["attention_mask"].to(device)
with torch.no_grad():
logits = model(input_ids, mask) # raw logits, shape (B, 17)
probs = torch.sigmoid(logits) # sigmoid -> probabilities
best_idx = torch.argmax(probs, dim=1) # index of highest prob per sample
return [labels[i] for i in best_idx.cpu().tolist()]
def predict(texts, bundle):
"""Route to the correct prediction function based on model type."""
if bundle.get("is_sdg", False):
return predict_with_sdg_model(texts, bundle)
else:
return predict_with_model(texts, bundle)
|