Update app.py
Browse files
app.py
CHANGED
|
@@ -51,6 +51,7 @@ def classify(text):
|
|
| 51 |
if not text or not text.strip():
|
| 52 |
return {}
|
| 53 |
enc = tokenizer(text, truncation=True, max_length=128, return_tensors="pt")
|
|
|
|
| 54 |
with torch.no_grad():
|
| 55 |
logits = model(**enc).logits
|
| 56 |
probs = F.softmax(logits, dim=-1)[0].tolist()
|
|
|
|
| 51 |
if not text or not text.strip():
|
| 52 |
return {}
|
| 53 |
enc = tokenizer(text, truncation=True, max_length=128, return_tensors="pt")
|
| 54 |
+
enc.pop("token_type_ids", None) # DistilBERT does not accept this argument
|
| 55 |
with torch.no_grad():
|
| 56 |
logits = model(**enc).logits
|
| 57 |
probs = F.softmax(logits, dim=-1)[0].tolist()
|