Visual LLM support (#341)

* visuallm with scripts

* corrected focus bug

* refactored visual lm to store images in dataset only

* small fixes

* fix

* corrected attentionscore

* fix

* fix

* fix transformers and default calculators

* Fix visual

* turn off focus

* fix attn

* fix attn

* fix attn

* get_images moving

* linter issues

* downgraded transformers

* linter issues

* linter issues

* linter issues

---------

Co-authored-by: alfekka <[email protected]>
Co-authored-by: iinemo <[email protected]>

Author: 3 people

examples/basic_example_visual.ipynb

@@ -0,0 +1,90 @@
+- name: MaximumSequenceProbability
+- name: Perplexity
+- name: MeanTokenEntropy
+- name: MeanPointwiseMutualInformation
+- name: MeanConditionalPointwiseMutualInformation 
+- name: ClaimConditionedProbability
+- name: PTrue
+- name: PTrueSampling
+- name: MonteCarloSequenceEntropy
+- name: MonteCarloNormalizedSequenceEntropy
+- name: LexicalSimilarity
+  cfg:
+    metric: "rouge1"
+- name: LexicalSimilarity
+  cfg:
+    metric: "rouge2"
+- name: LexicalSimilarity
+  cfg:
+    metric: "rougeL"
+- name: LexicalSimilarity
+  cfg:
+    metric: "BLEU"
+- name: NumSemSets
+- name: EigValLaplacian
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: EigValLaplacian
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "contra"
+- name: EigValLaplacian
+  cfg:
+    similarity_score: "Jaccard_score"
+- name: DegMat
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: DegMat
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "contra"
+- name: DegMat
+  cfg:
+    similarity_score: "Jaccard_score"
+- name: Eccentricity
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: Eccentricity
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "contra"
+- name: Eccentricity
+  cfg:
+    similarity_score: "Jaccard_score"
+- name: SemanticEntropy
+- name: SAR
+- name: TokenSAR
+- name: SentenceSAR
+- name: LUQ
+- name: KernelLanguageEntropy
+- name: EigenScore
+- name: RenyiNeg
+- name: FisherRao
+- name: MahalanobisDistanceSeq
+- name: RelativeMahalanobisDistanceSeq
+- name: RDESeq
+- name: PPLMDSeq
+  cfg:
+    md_type: "MD"
+- name: PPLMDSeq
+  cfg:
+    md_type: "RMD"
+- name: AttentionScore
+  cfg:
+    layer: 16
+    gen_only: False
+# - name: Focus
+#   cfg:
+#     model_name: '${model.path}'
+#     path: "${cache_path}/focus/${model.path}/token_idf.pkl"
+#     gamma: 0.9
+#     p: 0.01
+#     idf_dataset: "togethercomputer/RedPajama-Data-1T-Sample"
+#     trust_remote_code: True
+#     idf_seed: 42
+#     idf_dataset_size: 1000
+#     #idf_dataset_size: -1
+#     spacy_path: "en_core_web_sm"

examples/basic_visual_llm_example.ipynb

@@ -0,0 +1,22 @@
+from transformers import AutoModelForVision2Seq, AutoProcessor
+
+
+def load_model(model_path: str, device_map: str):
+    model = AutoModelForVision2Seq.from_pretrained(
+        model_path, trust_remote_code=True, device_map=device_map
+    )
+    model.eval()
+
+    return model
+
+
+def load_tokenizer(model_path: str):
+    processor_visual = AutoProcessor.from_pretrained(
+        model_path,
+        padding_side="left",
+        add_bos_token=True,
+    )
+    if processor_visual.tokenizer.pad_token is None:
+        processor_visual.tokenizer.pad_token = processor_visual.tokenizer.eos_token
+
+    return processor_visual

examples/configs/estimators/default_estimators_visual.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - default
+
+path: microsoft/kosmos-2-patch14-224
+type: VisualLM
+path_to_load_script: model/default_visual.py
+
+load_model_args:
+  device_map: auto
+load_tokenizer_args: {}

examples/configs/model/default_visual.py

@@ -0,0 +1,38 @@
+hydra:
+  run:
+    dir: ${cache_path}/${task}/${model}/${dataset}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+defaults:
+  - model: kosmos
+  - estimators: default_estimators_visual
+  - stat_calculators: default_calculators_visual
+  - _self_
+
+cache_path: ./workdir/output
+save_path: '${hydra:run.dir}'
+
+task: qa
+
+dataset: ['LM-Polygraph/vqa']
+text_column: question
+label_column: answer
+im_column: image
+train_split: train
+eval_split: test
+max_new_tokens: 3
+load_from_disk: false
+size: 100
+generation_params:
+  generate_until:
+    - "\n"
+
+subsample_eval_dataset: -1
+
+generation_metrics: null
+
+ignore_exceptions: false
+
+batch_size: 1
+
+seed:
+    - 1

examples/configs/model/kosmos.yaml

@@ -0,0 +1,32 @@
+- auto
+- name: TrainingStatisticExtractionCalculatorVisual
+  builder: lm_polygraph.defaults.stat_calculator_builders.default_TrainingStatisticExtractionCalculatorVisual
+  cfg:
+    dataset: '${dataset}'
+    text_column: '${text_column}'
+    label_column: '${label_column}'
+    im_column: '${im_column}'
+    description: ''
+    prompt: ''
+    few_shot_split: "train"
+    train_split: '${train_split}'
+    load_from_disk: '${load_from_disk}'
+    subsample_train_dataset: 50
+    n_shot: 5
+    background_train_dataset: LM-Polygraph/laion-1000-background
+    background_train_dataset_text_column: txt
+    background_train_dataset_label_column: __url__
+    background_load_from_disk: false
+    background_images: jpg
+    background_train_dataset_data_files: data/train-00000-of-00001.parquet
+    subsample_background_train_dataset: 100
+    batch_size: '${batch_size}'
+    seed: '${seed}'
+    size: '${size}'
+    bg_size: 1000
+    output_attentions: True
+  stats: 
+    - "train_embeddings"
+    - "background_train_embeddings"
+    - "train_greedy_log_likelihoods"
+  dependencies:

examples/configs/polygraph_eval_vqa.yaml

@@ -5,9 +5,9 @@ scikit-learn>=1.5.1
 tqdm>=4.64.1
 matplotlib>=3.6
 pandas>=1.3.5
-torch>=1.13.0
+torch>=2.6.0
 bs4
-transformers>=4.48.0,<4.52.0
+transformers==4.50.0
 nltk>=3.6.5
 sacrebleu>=1.5.0
 sentencepiece>=0.1.97

examples/configs/stat_calculators/default_calculators_visual.yaml

@@ -16,11 +16,12 @@ from lm_polygraph.utils.manager import UEManager
 from lm_polygraph.utils.dataset import Dataset
 from lm_polygraph.utils.model import WhiteboxModel, BlackboxModel
 from lm_polygraph.model_adapters import WhiteboxModelvLLM
+from lm_polygraph.model_adapters.visual_whitebox_model import VisualWhiteboxModel
 from lm_polygraph.utils.processor import Logger
 from lm_polygraph.generation_metrics import *
 from lm_polygraph.estimators import *
 from lm_polygraph.ue_metrics import *
-from lm_polygraph.utils.common import load_external_module
+from lm_polygraph.utils.common import load_external_module, load_processor, load_image
 from lm_polygraph.utils.generation_parameters import GenerationParameters, GenerationParametersFactory
 from lm_polygraph.defaults.register_default_stat_calculators import (
     register_default_stat_calculators,
@@ -30,6 +31,7 @@ from lm_polygraph.utils.builder_enviroment_stat_calculator import (
 )
 from lm_polygraph.utils.factory_estimator import FactoryEstimator
 from lm_polygraph.utils.factory_stat_calculator import StatCalculatorContainer
+#from transformers import AutoProcessor, AutoModelForVision2Seq
 
 hydra_config = Path(os.environ.get("HYDRA_CONFIG", ""))
 
@@ -115,12 +117,14 @@ def main(args):
             n_shot=getattr(args, "n_shot", 5),
             few_shot_split=getattr(args, "few_shot_split", "train"),
             few_shot_prompt=getattr(args, "few_shot_prompt", None),
+            im_column=getattr(args, "im_column", None),
             instruct=getattr(args, "instruct", None),
             split=args.eval_split,
             load_from_disk=args.load_from_disk,
             trust_remote_code=getattr(args, "trust_remote_code", False),
             **cache_kwargs,
         )
+#	images=dataset.images
         log.info("Done with loading eval data.")
 
         log.info("=" * 100)
@@ -191,7 +195,12 @@ def get_ue_metrics(args):
 
 
 def get_stat_calculator_names(config):
-    model_type = "Whitebox" if getattr(config.model, "type", "Whitebox") != "Blackbox" else "Blackbox"
+    model_type_raw = getattr(config.model, "type", "Whitebox")
+    model_type = (
+        "Blackbox" if model_type_raw == "Blackbox"
+        else "VisualLM" if model_type_raw == "VisualLM"
+        else "Whitebox"
+    )
     language = getattr(config, "language", "en")
     output_attentions = getattr(config, "output_attentions", True) and (getattr(config.model, "type", "Whitebox") != "vLLMCausalLM")
     output_hidden_states = False if getattr(config.model, "type", "Whitebox") == "vLLMCausalLM" else True
@@ -321,6 +330,12 @@ def get_generation_metrics(args):
 def get_model(args):
     if getattr(args.model, "type", "Whitebox") == "Blackbox":
         return get_blackbox_model(args)
+    elif getattr(args.model, "type", "Whitebox") == "VisualLM":
+        cache_kwargs = {
+            "cache_dir": getattr(args, "hf_cache", None),
+            "token": getattr(args, "hf_token", None),
+        }
+        return get_visual_model(args, cache_kwargs)
     elif getattr(args.model, "type", "Whitebox") == "vLLMCausalLM":
         return get_vllm_model(args)
     else:
@@ -404,6 +419,52 @@ def get_whitebox_model(args, cache_kwargs={}):
 
     return model
 
+
+def get_visual_model(args, cache_kwargs={}):
+    if not "path_to_load_script" in args.model or not args.model.path_to_load_script:
+        log.warning(
+            "Loading model by directly passing the path to the model is deprecated and will be removed in the next release. Please use loading script instead."
+        )
+        log.info(f"Loading model with cache_kwargs: {cache_kwargs}")
+        return VisualWhiteboxModel.from_pretrained(
+            args.model.path,
+            getattr(args, "generation_params", {}),
+            device_map=args.model.load_model_args.device_map,
+            add_bos_token=getattr(args.model, "add_bos_token", True),
+            **cache_kwargs
+        )
+
+    path_to_load_script = get_abs_path_from_hydra_config(
+            args.model.path_to_load_script
+        )
+    load_module = load_external_module(path_to_load_script)
+
+    load_model_args = {'model_path': args.model.path}
+    load_model_args.update(args.model.load_model_args)
+    base_model = load_module.load_model(**load_model_args)
+
+    load_tok_args = {'model_path': args.model.path}
+    load_tok_args.update(args.model.load_tokenizer_args)
+    tokenizer = load_module.load_tokenizer(**load_tok_args)
+
+    load_proc_args = {'model_path': args.model.path}
+    load_proc_args.update(getattr(args.model, "load_processor_args", {}))
+    processor = load_processor(**load_proc_args)
+
+    generation_params = GenerationParametersFactory.from_params(
+        yaml_config=getattr(args, "generation_params", {}),
+        native_config=base_model.generation_config.to_dict()
+    )
+
+    model = VisualWhiteboxModel(base_model,
+                          processor,
+                          args.model.path,
+                          args.model.type,
+                          generation_params)
+
+    return model
+
+
 def get_vllm_model(args):
     path_to_load_script = get_abs_path_from_hydra_config(
             args.model.path_to_load_script

requirements.txt

@@ -153,7 +153,8 @@ def _register(
             GreedyProbsVisualCalculator,
             "lm_polygraph.defaults.stat_calculator_builders.default_GreedyProbsVisualCalculator",
             {
-                "output_attentions": True,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
             },
         )
         _register(EntropyCalculator)
@@ -163,7 +164,6 @@ def _register(
         _register(BartScoreCalculator)
         _register(ModelScoreCalculator)
         _register(EnsembleTokenLevelDataCalculator)
-        _register(PromptVisualCalculator)
         _register(SamplingPromptVisualCalculator)
         _register(ClaimPromptVisualCalculator)
         _register(
@@ -193,6 +193,7 @@ def _register(
                 "language": language,
             },
         )
+        _register(AttentionForwardPassCalculatorVisual)
 
     else:
         raise NotImplementedError(f"Unknown model type: {model_type}")