Балаганский Никита Николаевич commited on
Commit
570bb74
1 Parent(s): aaf61e5
Files changed (1) hide show
  1. sampling.py +3 -0
sampling.py CHANGED
@@ -1,5 +1,7 @@
1
  from typing import List
2
 
 
 
3
  import torch
4
  from torch.nn import functional as F
5
 
@@ -28,6 +30,7 @@ class TopKWithTemperatureSampler:
28
 
29
 
30
  class CAIFSampler:
 
31
  def __init__(self, classifier_name, lm_tokenizer, device, invert_cls_probs: bool = False):
32
  self.device = device
33
  self.classifier_tokenizer = transformers.AutoTokenizer.from_pretrained(
 
1
  from typing import List
2
 
3
+ from functools import lru_cache
4
+
5
  import torch
6
  from torch.nn import functional as F
7
 
 
30
 
31
 
32
  class CAIFSampler:
33
+ @lru_cache(20)
34
  def __init__(self, classifier_name, lm_tokenizer, device, invert_cls_probs: bool = False):
35
  self.device = device
36
  self.classifier_tokenizer = transformers.AutoTokenizer.from_pretrained(