Балаганский Никита Николаевич
commited on
Commit
•
570bb74
1
Parent(s):
aaf61e5
add cache
Browse files- 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(
|