vortexSetupScripts/prepare_FLMR.py at main · aliciayuting/vortexSetupScripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from flmr import (
    FLMRConfig,
    FLMRVisionModel,
    FLMRContextEncoderTokenizer,
    FLMRModelForRetrieval,
    FLMRTextModel,
    FLMRQueryEncoderTokenizer,
    FLMRContextEncoderTokenizer, FLMRModelForRetrieval,
    )
from transformers import AutoImageProcessor


if __name__ == "__main__":

    checkpoint_path = 'LinWeizheDragon/PreFLMR_ViT-L'
    flmr_config = FLMRConfig.from_pretrained(checkpoint_path)
    query_tokenizer = FLMRQueryEncoderTokenizer.from_pretrained(
                    checkpoint_path,
                    text_config=flmr_config.text_config,
                    subfolder="query_tokenizer")
    context_tokenizer = FLMRContextEncoderTokenizer.from_pretrained(
                    checkpoint_path,
                    text_config=flmr_config.text_config,
                    subfolder="context_tokenizer")
    image_processor = AutoImageProcessor.from_pretrained('openai/clip-vit-large-patch14')
    flmr_model = FLMRModelForRetrieval.from_pretrained(
         checkpoint_path,
         query_tokenizer = query_tokenizer,
         context_tokenizer = context_tokenizer,
         )