Fix colab link

README.md

@@ -12,7 +12,7 @@ license: cc-by-nc-sa-4.0
 ---
 
 # <font color="IndianRed"> OTAS (Office Title Address Splitter)</font>
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1UoG3QebyBlK6diiYckiQv-5dRB9dA4iv?usp=sharing/)
+[![Open In Colab](https://colab.research.google.com/drive/1UoG3QebyBlK6diiYckiQv-5dRB9dA4iv?usp=sharing)
 
 Our model <font color="cornflowerblue">OTAS (Office Title Address Splitter) </font> is a Named Entity Recognition Classical Chinese language model that is intended to  <font color="IndianRed">split the address portion in Classical Chinese office titles.</font>. This model is first inherited from raynardj/classical-chinese-punctuation-guwen-biaodian Classical Chinese punctuation model, and finetuned using over a 25,000 high-quality punctuation pairs collected CBDB group (China Biographical Database). 
 
@@ -24,57 +24,69 @@ Here is how to use this model to get the features of a given text in PyTorch:
 ```python
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 
-device = torch.device('cuda')
-model_name = 'cbdb/OfficeTitleAddressSplitter'
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForTokenClassification.from_pretrained(model_name).to(device)
+PRETRAINED = "cbdb/OfficeTitleAddressSplitter"
+tokenizer = AutoTokenizer.from_pretrained(PRETRAINED)
+model = AutoModelForTokenClassification.from_pretrained(PRETRAINED)
 ```
 
 <font color="cornflowerblue"> 2. Load Data </font>
 ```python
 # Load your data here
-tobe_splitted = ['湖南常德協中軍都司','廣東鹽運使','漢軍鑲黃旗副都統']
+test_list = ['漢軍鑲黃旗副都統', '兵部右侍郎', '盛京戶部侍郎']
 ```
 
-work-in-progress
 
 <font color="cornflowerblue"> 3. Make a prediction </font>
 ```python
-
-  tokens_test = tokenizer.encode_plus(
-      tobe_splitted,
-      add_special_tokens=True,
-      return_attention_mask=True,
-      padding=True,
-      max_length=max_seq_len,
-      return_tensors='pt',
-      truncation=True
-  )
-
-  test_seq = torch.tensor(tokens_test['input_ids'])
-  test_mask = torch.tensor(tokens_test['attention_mask'])
-
-  # get predictions for test data
-  with torch.no_grad():
-    outputs = model(test_seq.cuda(), test_mask.cuda())
-    outputs = outputs.logits.detach().cpu().numpy()
-
-  softmax_score = softmax(outputs)
-  # pred_class_dict = {k:v for k, v in zip(label2idx.keys(), softmax_score[0])}
-  softmax_score = np.argmax(softmax_score, axis=2)[0]
-
-inputs = tokenizer(tobe_splitted, return_tensors="pt", padding=True).to(device)
-translated = model.generate(**inputs, max_length=128)
-tran = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
-for c, t in zip(tobe_translated, tran):
-    print(f'{c}: {t}')
+def predict_class(test):
+    tokens_test = tokenizer.encode_plus(
+        test,
+        add_special_tokens=True,
+        return_attention_mask=True,
+        padding=True,
+        max_length=128,
+        return_tensors='pt',
+        truncation=True
+    )
+
+    test_seq = torch.tensor(tokens_test['input_ids'])
+    test_mask = torch.tensor(tokens_test['attention_mask'])
+
+    inputs = {
+        "input_ids": test_seq,
+        "attention_mask": test_mask
+    }
+    with torch.no_grad():
+        # print(inputs.shape)
+        outputs = model(**inputs)
+        outputs = outputs.logits.detach().cpu().numpy()
+        
+    softmax_score = softmax(outputs)
+    softmax_score = np.argmax(softmax_score, axis=2)[0]
+    return test_seq, softmax_score
+
+for test_sen0 in test_list:
+    test_seq, pred_class_proba = predict_class(test_sen0)
+    test_sen = tokenizer.decode(test_seq[0]).split()
+    label = [idx2label[i] for i in pred_class_proba]
+
+    element_to_find = '。'
+
+    if element_to_find in label:
+        index = label.index(element_to_find)
+        test_sen_pred = [i for i in test_sen0]
+        test_sen_pred.insert(index, element_to_find)
+        test_sen_pred = ''.join(test_sen_pred)
+
+    else:
+        test_sen_pred = [i for i in test_sen0]
+        test_sen_pred = ''.join(test_sen_pred)
+
+    print(test_sen_pred)
 ```
-講筵官: Lecturer<br>
-判司簿尉: Supervisor of the Commandant of Records<br>
-散騎常侍: Policy Advisor<br>
-殿中省尚輦奉御: Chief Steward of the Palace Administration<br>
-
-work-in-progress
+漢軍鑲黃旗。副都統<br>
+兵部右侍郎<br>
+盛京。戶部侍郎<br>
 
 
 ### <font color="IndianRed">Authors </font>