Skip to content

Transformers

HFInference

Class for running HF model inference locally.

Source code in ragfit/models/hf.py 
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
class HFInference:
    """
    Class for running HF model inference locally.
    """

    def __init__(
        self,
        model_name_or_path: str,
        torch_dtype,
        device_map,
        instruction: Path,
        instruct_in_prompt: False,
        template: Path = None,
        lora_path=None,
        generation=None,
        task="text-generation",
        **kwargs,
    ):
        """
        Initialize a HF model, with optional LORA adapter.

        Args:
            model_name_or_path (str): HF model name or path.
            torch_dtype (str): torch dtype for the model.
            device_map: device map for the model.
            instruction (Path): path to the instruction file.
            instruct_in_prompt (bool): whether to include the instruction in the prompt for models without system role.
            template (Path): path to a prompt template file if tokenizer does not include chat template. Optional.
            lora_path (Path): path to the LORA adapter.
            generation (dict): generation kwargs.
            task (str): task for the pipeline.
        """

        self.model_name = model_name_or_path
        self.generation_kwargs = generation
        self.instruction = open(instruction).read()
        logger.info(f"Using the following instruction: {self.instruction}")

        self.instruct_in_prompt = instruct_in_prompt
        self.template = open(template).read() if template else None

        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, **kwargs)

        self.config = AutoConfig.from_pretrained(self.model_name, **kwargs)
        self.config.torch_dtype = torch_dtype or "auto"

        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name, config=self.config, device_map=device_map, **kwargs
        )
        if lora_path:
            logger.info(f"Loading LORA: {lora_path}")
            self.model.load_adapter(lora_path)

        self.pipe = pipeline(
            task=task,
            model=self.model,
            tokenizer=self.tokenizer,
        )

    def generate(self, prompt: str) -> str:
        """
        Given an input, generate a response.
        """

        if self.template:
            prompt = self.template.format(instruction=self.instruction, query=prompt)

        else:
            if self.instruct_in_prompt:
                prompt = self.instruction + "\n" + prompt

            messages = [
                {"role": "system", "content": self.instruction},
                {"role": "user", "content": prompt},
            ]

            prompt = self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                truncation=True,
                max_length=(
                    self.config.max_position_embeddings
                    - self.generation_kwargs["max_new_tokens"]
                ),
            )

        output = self.pipe(prompt, **self.generation_kwargs)
        return output[0]["generated_text"]

__init__(model_name_or_path: str, torch_dtype, device_map, instruction: Path, instruct_in_prompt: False, template: Path = None, lora_path=None, generation=None, task='text-generation', **kwargs)

Initialize a HF model, with optional LORA adapter.

Parameters:

  • model_name_or_path (str) –

    HF model name or path.

  • torch_dtype (str) –

    torch dtype for the model.

  • device_map –

    device map for the model.

  • instruction (Path) –

    path to the instruction file.

  • instruct_in_prompt (bool) –

    whether to include the instruction in the prompt for models without system role.

  • template (Path, default: None ) –

    path to a prompt template file if tokenizer does not include chat template. Optional.

  • lora_path (Path, default: None ) –

    path to the LORA adapter.

  • generation (dict, default: None ) –

    generation kwargs.

  • task (str, default: 'text-generation' ) –

    task for the pipeline.

Source code in ragfit/models/hf.py 
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
def __init__(
    self,
    model_name_or_path: str,
    torch_dtype,
    device_map,
    instruction: Path,
    instruct_in_prompt: False,
    template: Path = None,
    lora_path=None,
    generation=None,
    task="text-generation",
    **kwargs,
):
    """
    Initialize a HF model, with optional LORA adapter.

    Args:
        model_name_or_path (str): HF model name or path.
        torch_dtype (str): torch dtype for the model.
        device_map: device map for the model.
        instruction (Path): path to the instruction file.
        instruct_in_prompt (bool): whether to include the instruction in the prompt for models without system role.
        template (Path): path to a prompt template file if tokenizer does not include chat template. Optional.
        lora_path (Path): path to the LORA adapter.
        generation (dict): generation kwargs.
        task (str): task for the pipeline.
    """

    self.model_name = model_name_or_path
    self.generation_kwargs = generation
    self.instruction = open(instruction).read()
    logger.info(f"Using the following instruction: {self.instruction}")

    self.instruct_in_prompt = instruct_in_prompt
    self.template = open(template).read() if template else None

    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, **kwargs)

    self.config = AutoConfig.from_pretrained(self.model_name, **kwargs)
    self.config.torch_dtype = torch_dtype or "auto"

    self.model = AutoModelForCausalLM.from_pretrained(
        self.model_name, config=self.config, device_map=device_map, **kwargs
    )
    if lora_path:
        logger.info(f"Loading LORA: {lora_path}")
        self.model.load_adapter(lora_path)

    self.pipe = pipeline(
        task=task,
        model=self.model,
        tokenizer=self.tokenizer,
    )

generate(prompt: str) -> str

Given an input, generate a response.

Source code in ragfit/models/hf.py 
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
def generate(self, prompt: str) -> str:
    """
    Given an input, generate a response.
    """

    if self.template:
        prompt = self.template.format(instruction=self.instruction, query=prompt)

    else:
        if self.instruct_in_prompt:
            prompt = self.instruction + "\n" + prompt

        messages = [
            {"role": "system", "content": self.instruction},
            {"role": "user", "content": prompt},
        ]

        prompt = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            truncation=True,
            max_length=(
                self.config.max_position_embeddings
                - self.generation_kwargs["max_new_tokens"]
            ),
        )

    output = self.pipe(prompt, **self.generation_kwargs)
    return output[0]["generated_text"]

HFTrain

Class for training HF models locally.

Source code in ragfit/models/hf.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
class HFTrain:
    """
    Class for training HF models locally.
    """

    def __init__(
        self,
        model_name_or_path,
        torch_dtype,
        device_map,
        lora: LoraConfig = None,
        generation=None,
        completion_start: str = "",
        instruction_in_prompt=None,
        max_sequence_len=None,
        **kwargs,
    ):
        """
        Args:
            model_name_or_path: str - HF model name or path.
            torch_dtype: str - torch dtype for the model.
            device_map: dict - device map for the model.
            lora: dict - LoRA adapter config.
            generation: dict - generation kwargs.
            completion_start: str - used to find the start of the completion in the prompt.
            instruction_in_prompt: bool - whether to include the instruction in the prompt for models without system role.
        """
        self.model_name = model_name_or_path
        self.complete_start = completion_start
        self.instructions_in_prompt = instruction_in_prompt
        self.generation_kwargs = generation

        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        self.config = AutoConfig.from_pretrained(self.model_name, **kwargs)
        self.config.torch_dtype = torch_dtype or "auto"

        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            config=self.config,
            device_map=device_map,
            **kwargs,
        )

        self.model.config.use_cache = False
        logger.info(f"Loaded model: {self.model}")

        logger.info(f"Initializing LORA based on {lora}")
        self.model = get_peft_model(self.model, LoraConfig(**lora))

__init__(model_name_or_path, torch_dtype, device_map, lora: LoraConfig = None, generation=None, completion_start: str = '', instruction_in_prompt=None, max_sequence_len=None, **kwargs)

Parameters:

  • model_name_or_path –

    str - HF model name or path.

  • torch_dtype –

    str - torch dtype for the model.

  • device_map –

    dict - device map for the model.

  • lora (LoraConfig, default: None ) –

    dict - LoRA adapter config.

  • generation –

    dict - generation kwargs.

  • completion_start (str, default: '' ) –

    str - used to find the start of the completion in the prompt.

  • instruction_in_prompt –

    bool - whether to include the instruction in the prompt for models without system role.

Source code in ragfit/models/hf.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def __init__(
    self,
    model_name_or_path,
    torch_dtype,
    device_map,
    lora: LoraConfig = None,
    generation=None,
    completion_start: str = "",
    instruction_in_prompt=None,
    max_sequence_len=None,
    **kwargs,
):
    """
    Args:
        model_name_or_path: str - HF model name or path.
        torch_dtype: str - torch dtype for the model.
        device_map: dict - device map for the model.
        lora: dict - LoRA adapter config.
        generation: dict - generation kwargs.
        completion_start: str - used to find the start of the completion in the prompt.
        instruction_in_prompt: bool - whether to include the instruction in the prompt for models without system role.
    """
    self.model_name = model_name_or_path
    self.complete_start = completion_start
    self.instructions_in_prompt = instruction_in_prompt
    self.generation_kwargs = generation

    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
    self.tokenizer.pad_token = self.tokenizer.eos_token

    self.config = AutoConfig.from_pretrained(self.model_name, **kwargs)
    self.config.torch_dtype = torch_dtype or "auto"

    self.model = AutoModelForCausalLM.from_pretrained(
        self.model_name,
        config=self.config,
        device_map=device_map,
        **kwargs,
    )

    self.model.config.use_cache = False
    logger.info(f"Loaded model: {self.model}")

    logger.info(f"Initializing LORA based on {lora}")
    self.model = get_peft_model(self.model, LoraConfig(**lora))