update

pages/summarization_playground.py

@@ -33,8 +33,7 @@ Back in Boston, Kidd is going to rely on Lively even more. He'll play close to 3
     random_label: ""
 }
 
-
-def get_model_batch_generation(model_name):
+def model_device_check(model_name):
     global __model_on_gpu__
 
     if __model_on_gpu__ != model_name:
@@ -47,21 +46,15 @@ def get_model_batch_generation(model_name):
         model[model_name] = Model(model_name)
         __model_on_gpu__ = model_name
 
-    return model[model_name]
 
+def get_model_batch_generation(model_name):
+    model_device_check(model_name)
 
-def generate_answer(sources, model_name, prompt):
-    global __model_on_gpu__
+    return model[model_name]
 
-    if __model_on_gpu__ != model_name:
-        if __model_on_gpu__:
-            logging.info(f"delete model {__model_on_gpu__}")
-            del model[__model_on_gpu__]
-            gc.collect()
-            torch.cuda.empty_cache()
 
-        model[model_name] = Model(model_name)
-        __model_on_gpu__ = model_name
+def generate_answer(sources, model_name, prompt):
+    model_device_check(model_name)
 
     content = prompt + '\n{' + sources + '}\n\nsummary:'