Update app.py

app.py

@@ -56,29 +56,6 @@ model = model.to(device)
 model.eval()
 
 # Inference function
-def get_word_classifications(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, _ = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
-    word_tags = []
-    current_word = ""
-    current_tag = ""
-    for token, tag in zip(tokens, tags[0]):
-        if token in ["<s>", "</s>"]:
-            continue
-        if token.startswith("▁"):
-            if current_word:
-                word_tags.append(str(current_tag))
-            current_word = token[1:] if token != "▁" else ""
-            current_tag = tag
-        else:
-            current_word += token
-    if current_word:
-        word_tags.append(str(current_tag))
-    return word_tags
 
 def get_word_probabilities(text):
     text = " ".join(text.split(" ")[:2048])
@@ -172,7 +149,7 @@ def setup_hf_dataset():
 
 # Main inference + logging function
 def infer_and_log(text_input):
-    word_tags = get_word_classifications(text_input)
+    word_tags = get_word_probabilities(text_input)
     timestamp = datetime.datetime.now().isoformat()
     submission_id = str(uuid.uuid4())