Add `accelerate`

translate.py

@@ -6,9 +6,15 @@ import numpy as np
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
+from accelerate import Accelerator
+accelerator = Accelerator()
+
 tokenizer = AutoTokenizer.from_pretrained("under-tree/transformer-en-ru")
 model = AutoModelForSeq2SeqLM.from_pretrained("under-tree/transformer-en-ru")
 
+device = accelerator.device
+
+model = accelerator.prepare(model)
 
 @dataclass
 class TranslationResult:
@@ -24,15 +30,15 @@ class TranslationResult:
 
 def translator_fn(input_text: str, k=10) -> TranslationResult:
     # Preprocess input
-    inputs = tokenizer(input_text, return_tensors="pt")
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
     input_tokens = tokenizer.batch_decode(inputs.input_ids[0])
-    input_special_mask = torch.tensor([1 if t in tokenizer.all_special_tokens else 0 for t in input_tokens])
+    input_special_mask = torch.tensor([1 if t in tokenizer.all_special_tokens else 0 for t in input_tokens]).to(device)
 
     # Generate output
     outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, output_attentions=True)
     output_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
     output_tokens = tokenizer.batch_decode(outputs.sequences[0])
-    output_special_mask = torch.tensor([1 if t in tokenizer.all_special_tokens else 0 for t in output_tokens])
+    output_special_mask = torch.tensor([1 if t in tokenizer.all_special_tokens else 0 for t in output_tokens]).to(device)
 
     # Get cross attention matrix
     cross_attention = torch.stack([torch.stack(t) for t in outputs.cross_attentions])
@@ -61,8 +67,8 @@ def translator_fn(input_text: str, k=10) -> TranslationResult:
     clean_output_tokens = [t for t, m in zip(output_tokens, output_special_mask) if m == 0]
     clean_input_tokens = [t for t, m in zip(input_tokens, input_special_mask) if m == 0]
     clean_attention_matrix = attention_matrix[:len_output, :len_input]  # for padding
-    clean_attention_matrix = np.delete(clean_attention_matrix, np.where(output_special_mask == 1), axis=0)
-    clean_attention_matrix = np.delete(clean_attention_matrix, np.where(input_special_mask == 1), axis=1)
+    clean_attention_matrix = np.delete(clean_attention_matrix, np.where(output_special_mask.detach().cpu().numpy() == 1), axis=0)
+    clean_attention_matrix = np.delete(clean_attention_matrix, np.where(input_special_mask.detach().cpu().numpy() == 1), axis=1)
 
     n_input = len(clean_input_tokens)
     n_output = len(clean_output_tokens)