45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
"""Alpaca format strategy."""
|
||
|
||
from typing import List
|
||
|
||
from pipeline.tokenize import AutoTokenizer
|
||
from pipeline.strategies.base import PromptStrategy
|
||
from pipeline.strategies.factory import StrategyFactory
|
||
|
||
|
||
@StrategyFactory.register("alpaca")
|
||
class AlpacaStrategy(PromptStrategy):
|
||
"""Alpaca format:"""
|
||
|
||
def __init__(
|
||
self,
|
||
tokenizer: AutoTokenizer,
|
||
instruction_start: str = "### Instruction:\n",
|
||
response_start: str = "### Response:\n",
|
||
response_suffix: str = "\n<｜end▁of▁sentence｜>",
|
||
):
|
||
super().__init__(tokenizer)
|
||
self.instruction_start = instruction_start
|
||
self.response_start = response_start
|
||
self.response_suffix = response_suffix
|
||
|
||
self._instruction_start_ids = self._encode_format(instruction_start)
|
||
self._separator_ids = self._encode_format("\n\n")
|
||
self._response_start_ids = self._encode_format(response_start)
|
||
self._response_suffix_ids = self._encode_format(response_suffix)
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "alpaca"
|
||
|
||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||
return (
|
||
self._instruction_start_ids
|
||
+ query_tokens
|
||
+ self._separator_ids
|
||
+ self._response_start_ids
|
||
)
|
||
|
||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||
return response_tokens + self._response_suffix_ids
|