From f96862eb0436083973d40dc52e4419735801cdf0 Mon Sep 17 00:00:00 2001 From: Petr Andriushchenko Date: Fri, 24 Mar 2023 11:14:55 +0100 Subject: [PATCH] remove lavis from utils.py --- misinformation/utils.py | 36 ++---------------------------------- 1 file changed, 2 insertions(+), 34 deletions(-) diff --git a/misinformation/utils.py b/misinformation/utils.py index 1862b1d..6604cf8 100644 --- a/misinformation/utils.py +++ b/misinformation/utils.py @@ -2,8 +2,7 @@ import glob import os from pandas import DataFrame import pooch -from torch import device, cuda -from lavis.models import load_model_and_preprocess + class DownloadResource: @@ -107,35 +106,4 @@ if __name__ == "__main__": outdict = {} outdict = append_data_to_dict(mydict) df = dump_df(outdict) - print(df.head(10)) - - -def load_model_base(): - summary_device = device("cuda" if cuda.is_available() else "cpu") - summary_model, summary_vis_processors, _ = load_model_and_preprocess( - name="blip_caption", - model_type="base_coco", - is_eval=True, - device=summary_device, - ) - return summary_model, summary_vis_processors - - -def load_model_large(): - summary_device = device("cuda" if cuda.is_available() else "cpu") - summary_model, summary_vis_processors, _ = load_model_and_preprocess( - name="blip_caption", - model_type="large_coco", - is_eval=True, - device=summary_device, - ) - return summary_model, summary_vis_processors - - -def load_model(model_type): - select_model = { - "base": load_model_base, - "large": load_model_large, - } - summary_model, summary_vis_processors = select_model[model_type]() - return summary_model, summary_vis_processors + print(df.head(10)) \ No newline at end of file