Fix double model load

generator.py

@@ -18,7 +18,7 @@ def get_access_token():
  return access_token
 
 
-# @st.cache(hash_funcs={_thread.RLock: lambda _: None}, suppress_st_warning=True, allow_output_mutation=True)
+@st.cache_resource
 def load_model(model_name):
  os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
  tokenizer = AutoTokenizer.from_pretrained(
@@ -67,25 +67,24 @@ class Generator:
  self.load()
 
  def load(self):
- if not self.model:
- print(f"Loading model {self.model_name}")
- self.tokenizer, self.model = load_model(self.model_name)
-
- for key in self.gen_kwargs:
- if key in self.model.config.__dict__:
- self.gen_kwargs[key] = self.model.config.__dict__[key]
- try:
- if self.task in self.model.config.task_specific_params:
- task_specific_params = self.model.config.task_specific_params[
- self.task
- ]
- if "prefix" in task_specific_params:
- self.prefix = task_specific_params["prefix"]
- for key in self.gen_kwargs:
- if key in task_specific_params:
- self.gen_kwargs[key] = task_specific_params[key]
- except TypeError:
- pass
+ print(f"Loading model {self.model_name}")
+ self.tokenizer, self.model = load_model(self.model_name)
+
+ for key in self.gen_kwargs:
+ if key in self.model.config.__dict__:
+ self.gen_kwargs[key] = self.model.config.__dict__[key]
+ try:
+ if self.task in self.model.config.task_specific_params:
+ task_specific_params = self.model.config.task_specific_params[
+ self.task
+ ]
+ if "prefix" in task_specific_params:
+ self.prefix = task_specific_params["prefix"]
+ for key in self.gen_kwargs:
+ if key in task_specific_params:
+ self.gen_kwargs[key] = task_specific_params[key]
+ except TypeError:
+ pass
 
  def generate(self, text: str, streamer=None, **generate_kwargs) -> (str, dict):
  # Replace two or more newlines with a single newline in text
@@ -148,7 +147,6 @@ class GeneratorFactory:
  # If the generator is not yet present, add it
  if not self.get_generator(model_name=model_name, task=task, desc=desc):
  g = Generator(model_name, task, desc, split_sentences)
- g.load()
  self.generators.append(g)
 
  def get_generator(self, **kwargs):