everything wrapped in cache

generator.py

@@ -6,6 +6,8 @@ import streamlit as st
 
 from plotly import graph_objects as go
 
+from utils import get_lm
+
 
 class Generator:
     def __init__(self, lm_model_name, device, entropy=None):
@@ -15,9 +17,7 @@ class Generator:
         self.tokenizer = transformers.AutoTokenizer.from_pretrained(
             lm_model_name
         )
-        self.lm = transformers.AutoModelForCausalLM.from_pretrained(
-            lm_model_name
-        ).to(device)
+        self.lm = get_lm(lm_model_name).to(device)
         self.lm.eval()
 
         self.lm.config.pad_token_id = self.lm.config.eos_token_id

sampling.py

@@ -7,6 +7,8 @@ from torch.nn import functional as F
 
 import transformers
 
+from utils import get_cls
+
 
 def sample_from_values(unscaled_probs, values):
     samples = torch.multinomial(unscaled_probs, 1)
@@ -37,9 +39,7 @@ class CAIFSampler:
             classifier_name
         )
         self.classifier_model = (
-            transformers.AutoModelForSequenceClassification.from_pretrained(
-                classifier_name
-            ).to(device)
+            get_cls(classifier_name).to(device)
         )
         self.classifier_model.eval()
         self.lm_tokenizer = lm_tokenizer

utils.py

@@ -0,0 +1,13 @@
+from functools import lru_cache
+
+from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification
+
+
+@lru_cache(3)
+def get_lm(lm_name):
+    return AutoModelForCausalLM.from_pretrained(lm_name)
+
+
+@lru_cache(3)
+def get_cls(cls_name):
+    return AutoModelForSequenceClassification.from_pretrained(cls_name)