File size: 2,103 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 45 46 47 48 49 50 51 52 53 54 55 56 57 |
"""paligemma.py.
File for providing the Paligemma model implementation.
"""
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from src.models.base import ModelBase
from src.models.config import Config
class PaligemmaModel(ModelBase):
"""PaligemmaModel model implementation."""
def __init__(self, config: Config) -> None:
"""Initialization of the paligemma model.
Args:
config (Config): Parsed config
"""
# initialize the parent class
super().__init__(config)
def _load_specific_model(self) -> None:
"""Overridden function to populate Paligemma model.
Huggingface token is required to get access to the model.
Replace <HUGGINGFACE_TOKEN> in configs/paligemma-3b.yaml file with you own hugging face security token.
Note: 'token' is a general Hugging Face Hub access token, not specific to PaliGemma.
It enables loading private models or authenticated access.
See: https://huggingface.co/docs/hub/en/security-tokens
"""
self.model = PaliGemmaForConditionalGeneration.from_pretrained(
self.model_path, **self.config.model
)
def _init_processor(self) -> None:
"""Initialize the Paligemma processor.
Huggingface token is required.
Replace <HUGGINGFACE_TOKEN> in configs/paligemma-3b.yaml file with you own hugging face security token.
Note: 'token' is a general Hugging Face Hub access token, not specific to PaliGemma.
It enables loading private models or authenticated access.
See: https://huggingface.co/docs/hub/en/security-tokens
"""
self.processor = AutoProcessor.from_pretrained(self.model_path, token=self.config.model['token'])
def _generate_prompt(self, prompt: str) -> str:
"""Generates the Paligemma model prompt which will not use the chat template.
Args:
prompt (str): The input prompt for the model.
Returns:
str: The prompt to return, set by the config.
"""
return prompt
|