Update README.md

README.md

@@ -25,6 +25,8 @@ Model for extractive summarization based on [rubert-base-cased](DeepPavlov/ruber
 
 #### How to use
 
+Colab: [link](https://colab.research.google.com/drive/1Q8_v3H-kxdJhZIiyLYat7Kj02qDq7M1L)
+
 ```python
 import razdel
 from transformers import AutoTokenizer, BertForTokenClassification
@@ -48,7 +50,7 @@ inputs = tokenizer(
     truncation=True,
     return_tensors="pt",
 )
-sep_mask = inputs["input_ids"] == sep_token_id
+sep_mask = inputs["input_ids"][0] == sep_token_id
 
 # Fix token_type_ids
 current_token_type_id = 0 
@@ -60,7 +62,7 @@ for pos, input_id in enumerate(inputs["input_ids"][0]):
 # Infer model
 with torch.no_grad(): 
     outputs = model(**inputs) 
-logits = outputs.logits[:, :, 1][0]
+logits = outputs.logits[0, :, 1]
 
 # Choose sentences 
 logits = logits[sep_mask]
@@ -68,7 +70,7 @@ logits, indices = logits.sort(descending=True)
 logits, indices = logits.cpu().tolist(), indices.cpu().tolist()
 pairs = list(zip(logits, indices))
 pairs = pairs[:3]
-indices = [idx for _, idx in pairs] 
+indices = list(sorted([idx for _, idx in pairs]))
 summary = " ".join([sentences[idx] for idx in indices])
 print(summary)
 ```