Update app.py

app.py

@@ -674,16 +674,16 @@ def tag_SSCCV_span(sentence):
     tokens = SSCCVspanTaggingTokenizer.tokenize(
         SSCCVspanTaggingTokenizer.decode(SSCCVspanTaggingTokenizer.encode(sentence))
     )
-    inputs = SSCCVspanTaggingTokenizer.encode(sentence, return_tensors="pt")
-
-    outputs = SSCCVspanTaggingModel(inputs)[0]
+    inputs = SSCCVspanTaggingTokenizer.encode(sentence, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = SSCCVspanTaggingModel(inputs)[0]
     predictions = torch.argmax(outputs, dim=2)
 
     SSCCV_span_tagged_tokens = [
         (token, SSCCVspan_label_list[prediction])
         for token, prediction in zip(tokens, predictions[0].tolist())
     ]
-
+    
     SSCCV_span_tagged_predictions = [p for t, p in SSCCV_span_tagged_tokens]
 
     aligned_tokens = []
@@ -742,27 +742,42 @@ def tag_SSMN_span(sentence):
     if re.search(r"\[\:\]\s*$", sentence):
         sentence = sentence + " blah blah blah."
 
-    tokens = SSMNspanTaggingTokenizer.tokenize(
-        SSMNspanTaggingTokenizer.decode(SSMNspanTaggingTokenizer.encode(sentence))
+    # --- BERT side ---
+    # Prefer the full call so you get input_ids + attention_mask
+    encoded = SSMNspanTaggingTokenizer(
+        sentence,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512,
     )
-    inputs = SSMNspanTaggingTokenizer.encode(sentence, return_tensors="pt")
-    if device != "cpu":
-        inputs = inputs.to(device)
+
+    # Put every tensor on the *same* device as the model
+    model_device = next(SSMNspanTaggingModel.parameters()).device
+    encoded = {k: v.to(model_device) for k, v in encoded.items()}
 
     with torch.no_grad():
-        outputs = SSMNspanTaggingModel(inputs)[0]
+        outputs = SSMNspanTaggingModel(**encoded)[0]
     predictions = torch.argmax(outputs, dim=2)
 
+    # Wordpiece tokens for alignment (same string as above)
+    tokens = SSMNspanTaggingTokenizer.tokenize(
+        SSMNspanTaggingTokenizer.decode(
+            SSMNspanTaggingTokenizer.encode(sentence)
+        )
+    )
+
     ssmn_span_tagged_predictions = [
-        SSMNspan_label_list[prediction]
-        for prediction in predictions[0].tolist()
+        SSMNspan_label_list[p] for p in predictions[0].tolist()
     ]
 
+    # --- spaCy side (CPU is fine) ---
     aligned_tokens, aligned_predictions = [], []
     aligned_pos_tags, aligned_lemmas = [], []
 
     sc_sentence = (
-        clean_for_spaCy(sentence) if "[" in sentence and "]" in sentence else sentence
+        clean_for_spaCy(sentence)
+        if "[" in sentence and "]" in sentence
+        else sentence
     )
     doc = nlp(sc_sentence)
     spacy_pos_tags = [t.tag_ for t in doc]
@@ -783,6 +798,7 @@ def tag_SSMN_span(sentence):
     return sign_cleaning(
         aligned_tokens, aligned_pos_tags, aligned_lemmas, aligned_predictions
     )
+    
 ###############################################################################
 
 
@@ -974,31 +990,45 @@ def tag_CMN1_span(sentence):
 
 ###############################################################################
 def tag_CLN_span(sentence):
-    if re.search("\[\:\]\s*$", sentence):
+    if re.search(r"\[\:\]\s*$", sentence):
         sentence = sentence + " blah blah blah."
 
-    tokens = CLNspanTaggingTokenizer.tokenize(
-        CLNspanTaggingTokenizer.decode(CLNspanTaggingTokenizer.encode(sentence))
+    # --- BERT side ---
+    encoded = CLNspanTaggingTokenizer(
+        sentence,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512,
     )
-    inputs = CLNspanTaggingTokenizer.encode(sentence, return_tensors="pt")
 
-    outputs = CLNspanTaggingModel(inputs)[0]
+    # Same device as the model (works on Windows CPU and HF GPU/ZeroGPU)
+    model_device = next(CLNspanTaggingModel.parameters()).device
+    encoded = {k: v.to(model_device) for k, v in encoded.items()}
+
+    with torch.no_grad():
+        outputs = CLNspanTaggingModel(**encoded)[0]
     predictions = torch.argmax(outputs, dim=2)
 
-    CLN_span_tagged_tokens = [
-        (token, CLNspan_label_list[prediction])
-        for token, prediction in zip(tokens, predictions[0].tolist())
-    ]
+    tokens = CLNspanTaggingTokenizer.tokenize(
+        CLNspanTaggingTokenizer.decode(
+            CLNspanTaggingTokenizer.encode(sentence)
+        )
+    )
 
-    CLN_span_tagged_predictions = [p for t, p in CLN_span_tagged_tokens]
+    CLN_span_tagged_predictions = [
+        CLNspan_label_list[p] for p in predictions[0].tolist()
+    ]
 
+    # --- spaCy side (CPU) ---
     aligned_tokens = []
     aligned_predictions = []
     aligned_pos_tags = []
     aligned_lemmas = []
 
     sc_sentence = (
-        clean_for_spaCy(sentence) if "[" in sentence and "]" in sentence else sentence
+        clean_for_spaCy(sentence)
+        if "[" in sentence and "]" in sentence
+        else sentence
     )
     doc = nlp(sc_sentence)
     spacy_pos_tags = [t.tag_ for t in doc]
@@ -1009,10 +1039,11 @@ def tag_CLN_span(sentence):
 
     for i in range(len(spacy_tokens)):
         for j in a2b[i]:
-            aligned_tokens.append(spacy_tokens[i])
-            aligned_pos_tags.append(spacy_pos_tags[i])
-            aligned_lemmas.append(spacy_lemmas[i])
-            aligned_predictions.append(CLN_span_tagged_predictions[j])
+            if j < len(CLN_span_tagged_predictions):
+                aligned_tokens.append(spacy_tokens[i])
+                aligned_pos_tags.append(spacy_pos_tags[i])
+                aligned_lemmas.append(spacy_lemmas[i])
+                aligned_predictions.append(CLN_span_tagged_predictions[j])
             break
 
     clean_tokens, clean_pos_tags, clean_lemmas, clean_predictions = sign_cleaning(
@@ -1029,9 +1060,7 @@ def tag_CLN_span(sentence):
         clean_pos_tags,
         clean_lemmas,
         clean_predictions,
-    )  #   added return
-
-
+    )
 ###############################################################################
 
 ###############################################################################