File size: 3,056 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
"""internlm.py.
File for providing the InternLM-XComposer model implementation.
"""
import logging
import torch
from transformers import AutoModel, AutoProcessor
from src.models.base import ModelBase
from src.models.config import Config
class InternLMXComposerModel(ModelBase):
"""InternLM model implementation."""
def __init__(self, config: Config) -> None:
"""Initialization of the InternLM 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 = AutoModel.from_pretrained(
self.model_path,
trust_remote_code=True,
**self.config.model
) if hasattr(self.config, 'model') else (
AutoModel.from_pretrained(
self.model_path,
trust_remote_code=True
)
)
def _init_processor(self) -> None:
"""Overridden function to instantiate the model's processor."""
self.processor = AutoProcessor.from_pretrained(
self.model_path, trust_remote_code=True)
self.model.tokenizer = self.processor
def _generate_prompt(self, prompt: str) -> str:
"""Overridden function to generate the prompt for the model.
Args:
prompt (str): The input prompt to be processed.
Returns:
str: The formatted prompt ready for model input.
"""
return prompt
def _generate_processor_output(self, prompt: str, img_path: str) -> dict:
"""Overridden function to generate the format the prompt for the processor.
Args:
prompt (str): The input prompt to be processed.
img_path (str): The path to the image to be processed.
Returns:
dict: The formatted inputs for the processor.
Raises:
ValueError: If no prompt is provided when required.
"""
logging.debug('Loading data...')
# Manually format input as we do not need a processor
inputs = {}
# Text prompts are required for this model
if not prompt:
raise ValueError(
'No input prompt was provided for the InternLM-XC model')
# If there are images, load them and add image token to prompt
if self.config.has_images():
inputs['query'] = f'<ImageHere>; {prompt}'
inputs['image'] = [img_path]
else:
inputs['query'] = prompt
return inputs
def _forward(self, data: dict) -> None:
"""Overridden function to run the model forward pass.
Args:
data (dict): The input data for the model.
"""
device_type = str(self.config.device)
logging.debug(f'DATA: {data}')
with torch.autocast(device_type=device_type):
_, _ = self.model.chat(
self.processor, **data, **self.config.forward)
|