Add application file

app.py

@@ -1,5 +1,4 @@
 import streamlit as st
-import torch
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 import time
 import json
@@ -67,13 +66,9 @@ MODEL_DETAILS = {
 class PIIMaskingModel:
     def __init__(self, model_name: str):
         self.model_name = model_name
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
-        self.model = AutoModelForTokenClassification.from_pretrained(
-            model_name, token=HF_TOKEN
-        )
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
-        self.model.eval()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForTokenClassification.from_pretrained(model_name)
+        # No need to specify device as we are forcing CPU usage
 
     def process_text(
         self, text: str
@@ -84,23 +79,23 @@ class PIIMaskingModel:
             text,
             truncation=True,
             padding=False,
-            return_tensors="pt",
+            return_tensors="np",  # Return NumPy arrays for CPU
             return_offsets_mapping=True,
             add_special_tokens=True,
         )
 
-        input_ids = tokenized_inputs.input_ids.to(self.device)
-        attention_mask = tokenized_inputs.attention_mask.to(self.device)
+        input_ids = tokenized_inputs.input_ids
+        attention_mask = tokenized_inputs.attention_mask
         offset_mapping = tokenized_inputs["offset_mapping"][0].tolist()
 
         # Handle special tokens
         offset_mapping[0] = None  # <s> token
         offset_mapping[-1] = None  # </s> token
 
-        with torch.no_grad():
-            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        # No need for torch.no_grad() as we are not using gradients
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
 
-        predictions = outputs.logits.argmax(dim=-1).cpu().numpy()
+        predictions = outputs.logits.argmax(dim=-1)  # No need to move to CPU
         predicted_labels = [
             self.model.config.id2label[label_id] for label_id in predictions[0]
         ]
@@ -140,7 +135,7 @@ class PIIMaskingModel:
             next_label = labels[j]
 
             # Stop if we hit a new B- tag (except for non-spaced tokens)
-            if next_label.startswith("B-") and tokens[j].startswith("▁"):
+            if next_label.startswith("B-") and tokens[j].startswith(" "):
                 break
 
             # Stop if we hit a different entity type in I- tags
@@ -152,7 +147,7 @@ class PIIMaskingModel:
                 last_valid_end = offset_mapping[j][1]
                 j += 1
             # Continue if it's a non-spaced B- token
-            elif next_label.startswith("B-") and not tokens[j].startswith("▁"):
+            elif next_label.startswith("B-") and not tokens[j].startswith(" "):
                 last_valid_end = offset_mapping[j][1]
                 j += 1
             else: