Update deploy_pipeline.py

deploy_pipeline.py

@@ -1,14 +1,3 @@
-import torch
-from transformers import Pipeline
-from transformers import AutoTokenizer
-from transformers.pipelines import PIPELINE_REGISTRY
-from transformers import pipeline
-from transformers import AutoModelForTokenClassification
-from huggingface_hub import Repository
-import sys
-import os
-
-
 class TokenizeAndAlignLabelsStep():
 
     # Adapted From : https://huggingface.co/docs/transformers/tasks/token_classification
@@ -21,7 +10,7 @@ class TokenizeAndAlignLabelsStep():
 
         previous_word_idx = None
         
-    
+        tokens= []
         labels_mask = []        
         
         for word_idx in word_ids:  # Set the special tokens to -100.
@@ -29,64 +18,17 @@ class TokenizeAndAlignLabelsStep():
                 labels_mask.append(False)
             # Only label the first token of a given word.
             elif word_idx != previous_word_idx:                    
-                labels_mask.append(True)                
+                labels_mask.append(True)
+                tokens.append(tokenized_inputs["input_ids"][word_idx])
             else:                    
                 labels_mask.append(False)
             
             previous_word_idx = word_idx
                         
-        tokenized_inputs["tokens"] = tokenizer.decode(tokenized_inputs["input_ids"], skip_special_tokens=True)
+        tokenized_inputs["tokens"] = tokens
         tokenized_inputs["labels_mask"] = labels_mask
 
         return tokenized_inputs
-
-
-class BERT_CRF_Pipeline(Pipeline):
-
-    def _sanitize_parameters(self, **kwargs):
-        return {}, {}, {}
-
-    def preprocess(self, text):
-
-        tokenizer = AutoTokenizer.from_pretrained(
-            "neuralmind/bert-base-portuguese-cased", do_lower_case=False)            
-
-        TokenizeAndAlignLabelsStep().tokenize_and_align_labels(
-            examples=text, tokenizer=tokenizer)
-
-        return TokenizeAndAlignLabelsStep().tokenize_and_align_labels(examples=text, tokenizer=tokenizer)
-
-    def _forward(self, tokenizer_results):
-
-        input_ids = torch.tensor(
-            tokenizer_results['input_ids'], dtype=torch.long, device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")).unsqueeze(0)
-
-        token_type_ids = torch.tensor(
-            tokenizer_results['token_type_ids'], dtype=torch.long, device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")).unsqueeze(0)
-
-        attention_mask = torch.tensor(
-            tokenizer_results['attention_mask'], dtype=torch.bool, device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")).unsqueeze(0)
-
-        labels_mask = torch.tensor(
-            tokenizer_results['labels_mask'], dtype=torch.bool, device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")).unsqueeze(0)
-
-        # input_ids, token_type_ids, attention_mask, labels, labels_mask
-        outputs = self.model(input_ids=input_ids, token_type_ids=token_type_ids,
-                             attention_mask=attention_mask, labels=None, labels_mask=labels_mask)
-
-        return {'outputs': outputs, 'tokens': tokenizer_results['tokens']}
-
-    def postprocess(self, outputs):
-        
-        model_outputs = outputs['outputs']
-        tokens = outputs['tokens']
-
-        # From Ner_tags to Ner_labels
-        for i, label in enumerate(model_outputs[0]):
-            model_outputs[0][i] = self.model.config.id2label[label]
-                
-        return model_outputs[0], tokens
-
 def main():    
 
     PIPELINE_REGISTRY.register_pipeline("PT-BERT-Large-CRF-HAREM-Default-pipeline",