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):