readme: ---
  tags:
  - ner
  - punctuation
  language:
  - zh
  ---
  # zh-wiki-punctuation-restore
  More Detail: https://github.com/p208p2002/ZH-Punctuation-Restore
  
  共計支援6種標點符號: ， 、 。 ？ ！ ； 
  
  ## Install
  ```bash
  # pip install torch pytorch-lightning
  pip install zhpr
  ```
  
  ## Usage
  ```python
  from zhpr.predict import DocumentDataset,merge_stride,decode_pred
  from transformers import AutoModelForTokenClassification,AutoTokenizer
  from torch.utils.data import DataLoader
  
  def predict_step(batch,model,tokenizer):
          batch_out = []
          batch_input_ids = batch
  
          encodings = {'input_ids': batch_input_ids}
          output = model(**encodings)
  
          predicted_token_class_id_batch = output['logits'].argmax(-1)
          for predicted_token_class_ids, input_ids in zip(predicted_token_class_id_batch, batch_input_ids):
              out=[]
              tokens = tokenizer.convert_ids_to_tokens(input_ids)
              
              # compute the pad start in input_ids
              # and also truncate the predict
              # print(tokenizer.decode(batch_input_ids))
              input_ids = input_ids.tolist()
              try:
                  input_id_pad_start = input_ids.index(tokenizer.pad_token_id)
              except:
                  input_id_pad_start = len(input_ids)
              input_ids = input_ids[:input_id_pad_start]
              tokens = tokens[:input_id_pad_start]
      
              # predicted_token_class_ids
              predicted_tokens_classes = [model.config.id2label[t.item()] for t in predicted_token_class_ids]
              predicted_tokens_classes = predicted_tokens_classes[:input_id_pad_start]
  
              for token,ner in zip(tokens,predicted_tokens_classes):
                  out.append((token,ner))
              batch_out.append(out)
          return batch_out
  
  if __name__ == "__main__":
      window_size = 256
      step = 200
      text = "維基百科是維基媒體基金會運營的一個多語言的百科全書目前是全球網路上最大且最受大眾歡迎的參考工具書名列全球二十大最受歡迎的網站特點是自由內容自由編輯與自由著作權"
      dataset = DocumentDataset(text,window_size=window_size,step=step)
      dataloader = DataLoader(dataset=dataset,shuffle=False,batch_size=5)
  
      model_name = 'p208p2002/zh-wiki-punctuation-restore'
      model = AutoModelForTokenClassification.from_pretrained(model_name)
      tokenizer = AutoTokenizer.from_pretrained(model_name)
  
      model_pred_out = []
      for batch in dataloader:
          batch_out = predict_step(batch,model,tokenizer)
          for out in batch_out:
              model_pred_out.append(out)
          
      merge_pred_result = merge_stride(model_pred_out,step)
      merge_pred_result_deocde = decode_pred(merge_pred_result)
      merge_pred_result_deocde = ''.join(merge_pred_result_deocde)
      print(merge_pred_result_deocde)
  ```
  ```
  維基百科是維基媒體基金會運營的一個多語言的百科全書，目前是全球網路上最大且最受大眾歡迎的參考工具書，名列全球二十大最受歡迎的網站，特點是自由內容、自由編輯與自由著作權。
  ```