File size: 1,152 Bytes
d425e71
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
32
33
34
35
36
37
38
39
40
41
42
43
44
"""clip.py.

File for providing the Clip model implementation.
"""

from transformers import CLIPModel

from src.models.base import ModelBase
from src.models.config import Config


class ClipModel(ModelBase):
    """Clip model implementation."""

    def __init__(self, config: Config) -> None:
        """Initialization of the clip model.

        Args:
            config (Config): Parsed config
        """
        # initialize the parent class
        super().__init__(config)

    def _load_specific_model(self) -> None:
        """Overridden function to populate self.model."""
        self.model = CLIPModel.from_pretrained(
            self.model_path, **self.config.model
        ) if hasattr(self.config, 'model') else (
            CLIPModel.from_pretrained(
                self.model_path
            )
        )

    def _generate_prompt(self, prompt: str) -> str:
        """Generates the CLIP model prompt which will not use the chat template.

        Args:
            prompt (str): The prompt to return, set by the config.

        Returns:
            str: The prompt to return, set by the config.
        """
        return prompt