Update app.py

app.py

@@ -80,6 +80,38 @@ def get_word_classifications(text):
         word_tags.append(str(current_tag))
     return word_tags
 
+def get_word_probabilities(text):
+    text = " ".join(text.split(" ")[:2048])
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    with torch.no_grad():
+        tags, emission = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
+    probs = torch.softmax(emission, dim=-1)[0, :, 1].cpu().numpy()
+        
+    word_probs = []
+    current_word = ""
+    current_probs = []
+
+    for token, prob in zip(tokens, probs):
+        if token in ["<s>", "</s>"]:
+            continue
+        if token.startswith("▁"):
+            if current_word and current_probs:
+                word_probs.append(sum(current_probs) / len(current_probs))
+            current_word = token[1:] if token != "▁" else ""
+            current_probs = [prob]
+        else:
+            current_word += token
+            current_probs.append(prob)
+
+    if current_word and current_probs:
+        word_probs.append(sum(current_probs) / len(current_probs))
+
+    return word_probs
+
+    
+
 # def get_word_classifications(text):
 #     text = " ".join(text.split(" ")[:2048])
 #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)