diff --git a/benchmarks/bench_processors.py b/benchmarks/bench_processors.py new file mode 100644 index 0000000000..96b74b78fd --- /dev/null +++ b/benchmarks/bench_processors.py @@ -0,0 +1,44 @@ +import mlx.core as mx +import numpy as np +import torch + +from outlines.processors import OutlinesLogitsProcessor + + +class HalvingLogitsProcessor(OutlinesLogitsProcessor): + """Simply halve the passed logits""" + + def process_logits(self, input_ids, logits): + return logits / 2 + + +class LogitsProcessorBenchmark: + params = ["torch", "numpy"] + if mx.metal.is_available(): + params += ["mlx"] + + def setup(self, array_library): + self.logits_processor = HalvingLogitsProcessor() + + # logits: (4, 30,000 ) dtype=float + # input_ids shape: (4, 2048) dtype=int + if array_library == "torch": + self.logits = torch.rand((4, 30000), dtype=torch.float) + self.input_ids = torch.randint( + low=0, high=30000, size=(4, 2048), dtype=torch.int + ) + elif array_library == "numpy": + self.logits = np.random.rand(4, 30000).astype(np.float32) + self.input_ids = np.random.randint(low=0, high=30000, size=(4, 2048)) + elif array_library == "mlx": + self.logits = mx.random.uniform( + low=-1e9, high=1e9, shape=(4, 30000), dtype=mx.float32 + ) + self.input_ids = mx.random.randint( + low=0, high=30000, shape=(4, 2048), dtype=mx.int32 + ) + else: + raise ValueError + + def time_logits_processor(self, array_library): + self.logits_processor(self.input_ids, self.logits) diff --git a/outlines/__init__.py b/outlines/__init__.py index 3eb6a2f943..307d2ba6f4 100644 --- a/outlines/__init__.py +++ b/outlines/__init__.py @@ -2,6 +2,7 @@ import outlines.generate import outlines.grammars import outlines.models +import outlines.processors import outlines.types from outlines.base import vectorize from outlines.caching import clear_cache, disable_cache, get_cache diff --git a/outlines/fsm/guide.py b/outlines/fsm/guide.py index d247db62be..f8077dd5b6 100644 --- a/outlines/fsm/guide.py +++ b/outlines/fsm/guide.py @@ -11,6 +11,7 @@ make_byte_level_fsm, make_deterministic_fsm, ) +from outlines.fsm.parsing import PartialLark, terminals_to_fsms if TYPE_CHECKING: from outlines.models.tokenizer import Tokenizer @@ -256,15 +257,23 @@ def __init__(self, cfg_string: str, tokenizer): self.cfg_string = cfg_string self.tokenizer = tokenizer - self.parser = Lark( + self.parser = PartialLark( cfg_string, parser="lalr", - lexer="contextual", - propagate_positions=False, - maybe_placeholders=False, - regex=True, + deterministic=True, import_paths=[grammars.GRAMMAR_PATH], + # TODO: old options, not sure we need them, investigate + # propagate_positions=False, + # maybe_placeholders=False, + + # TODO: old PartialLark options, investigate + # start="file_input", ) + + self.regex_fsm = terminals_to_fsms(self.parser) + self.generation = "" + + """ self.terminal_regexps = dict() for terminal in self.parser.terminals: if terminal.pattern is not None: @@ -279,6 +288,7 @@ def __init__(self, cfg_string: str, tokenizer): self.check_last = False self.proposal_last: List[int] = [] self.regex_fsm_last: RegexGuide + """ self.start_state = 0 self.final_state = -1 @@ -316,6 +326,9 @@ def get_next_instruction(self, state: int) -> Instruction: A list that contains the tokens to mask. """ + + import pdb;pdb.set_trace() + if self.is_final_state(state): return Write([self.tokenizer.eos_token_id]) diff --git a/outlines/generate/cfg.py b/outlines/generate/cfg.py index e473c26a6c..a112af3a43 100644 --- a/outlines/generate/cfg.py +++ b/outlines/generate/cfg.py @@ -2,10 +2,7 @@ from outlines.fsm.guide import CFGGuide from outlines.generate.api import SequenceGenerator, SequenceGeneratorAdapter -from outlines.models import OpenAI -from outlines.models.llamacpp import LlamaCpp -from outlines.models.mlxlm import MLXLM -from outlines.models.vllm import VLLM +from outlines.models import MLXLM, VLLM, LlamaCpp, OpenAI, Transformers from outlines.samplers import Sampler, multinomial @@ -36,25 +33,16 @@ def cfg(model, cfg_str: str, sampler: Sampler = multinomial()) -> SequenceGenera @cfg.register(MLXLM) @cfg.register(VLLM) -def cfg_unimplemented( - model, - cfg_str: str, - sampler: Sampler = multinomial(), -): - raise NotImplementedError( - f"The CFG Logits processor is not available for {type(model)}." - ) - - @cfg.register(LlamaCpp) -def cfg_llamacpp( - model: LlamaCpp, +@cfg.register(Transformers) +def cfg_unified( + model, cfg_str: str, sampler: Sampler = multinomial(), ): - from outlines.integrations.llamacpp import CFGLogitsProcessor + from outlines.processors import CFGLogitsProcessor - logits_processor = CFGLogitsProcessor(cfg_str, model.model) + logits_processor = CFGLogitsProcessor(cfg_str, tokenizer=model.tokenizer) return SequenceGeneratorAdapter(model, logits_processor, sampler) diff --git a/outlines/grammars.py b/outlines/grammars.py index f0c1229647..c85635b714 100644 --- a/outlines/grammars.py +++ b/outlines/grammars.py @@ -12,3 +12,4 @@ def read_grammar(grammar_file_name, base_grammar_path=GRAMMAR_PATH): arithmetic = read_grammar("arithmetic.lark") json = read_grammar("json.lark") +sql_select = read_grammar("sql_select.lark") diff --git a/outlines/grammars/common.lark b/outlines/grammars/common.lark index 801c27e97d..3f1bd42d51 100644 --- a/outlines/grammars/common.lark +++ b/outlines/grammars/common.lark @@ -43,11 +43,12 @@ SIGNED_FLOAT: ["+"|"-"] FLOAT NUMBER: FLOAT | INT SIGNED_NUMBER: ["+"|"-"] NUMBER -// -// TODO: Working escaped_string -// UNESCAPED_STRING: /\"[^"]*\"/ +// based on `outlines/fsm/json_schema.py` +ESCAPED_STRING_INNER: /([^"\\\\\\x00-\\x1F\\x7F-\\x9F]|\\\\["\\\\])/ +ESCAPED_STRING: "\"" ESCAPED_STRING_INNER* "\"" + // diff --git a/outlines/grammars/sql_select.lark b/outlines/grammars/sql_select.lark new file mode 100644 index 0000000000..72f7d1fd18 --- /dev/null +++ b/outlines/grammars/sql_select.lark @@ -0,0 +1,203 @@ +// Adapted from https://github.com/zbrookle/sql_to_ibis +// License for https://github.com/zbrookle/sql_to_ibis follows +//BSD 3-Clause License +// +//Copyright (c) 2011-2022, Open source contributors. +// +//Redistribution and use in source and binary forms, with or without +//modification, are permitted provided that the following conditions are met: +// +//* Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +//* Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +//* Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +//THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +//AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +//IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +//DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +//FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +//DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +//SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +//CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +//OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +//OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +start: set_expr -> final + +set_expr: query_expr + | set_expr "UNION"i ["DISTINCT"i] set_expr -> union_distinct + | set_expr "UNION"i "ALL"i set_expr -> union_all + | set_expr "INTERSECT"i ["DISTINCT"i] set_expr -> intersect_distinct + | set_expr "EXCEPT"i ["DISTINCT"i] set_expr -> except_distinct + | set_expr "EXCEPT"i "ALL"i set_expr -> except_all + +query_expr: select [ "ORDER"i "BY"i (order_by_expr ",")* order_by_expr] [ "LIMIT"i limit_count [ "OFFSET"i skip_rows ] ] + +select: "SELECT"i [SELECT_CONSTRAINT] [(select_expr ",")*] select_expr "FROM"i [(from_expr ",")*] from_expr [ "WHERE"i where_expr ] [ "GROUP"i "BY"i [(groupby_expr ",")*] groupby_expr ] [ "HAVING"i having_expr] [ "WINDOW"i window_expr ] + +where_expr: bool_expression + +select_expr.0: expression_math [ [ "AS"i ] alias ] -> select_expression + +?from_expr: from_item -> from_expression + +order_by_expr: order -> order_by_expression + +having_expr: bool_expression + +groupby_expr: expression -> group_by + +window_expr: [window_expr ","] _window_name "AS"i ( window_definition ) + +from_item: name [ [ "AS"i ] alias ] -> table + | join -> join + | cross_join -> cross_join_expression + | subquery + +subquery: ( "(" (query_expr | join | cross_join) ")" ) [ [ "AS"i ] alias ] + +cross_join: from_item "CROSS"i "JOIN"i from_item +join: from_item JOIN_EXPR from_item [ "ON"i bool_expression ] -> join_expression + +JOIN_EXPR.5: (JOIN_TYPE WS)? "JOIN"i +JOIN_TYPE: "INNER"i | "OUTER"i? | JOIN_DIRECTION (WS "OUTER"i)? | JOIN_DIRECTION +JOIN_DIRECTION: "FULL"i | "LEFT"i | "RIGHT"i + +?expression_math: expression_product + | expression_math "+" expression_product -> expression_add + | expression_math "-" expression_product -> expression_sub + | "CASE"i (when_then)+ "ELSE"i expression_math "END"i -> case_expression + | "CAST"i "(" expression_math "AS"i TYPENAME ")" -> as_type + | "CAST"i "(" literal "AS"i TYPENAME ")" -> literal_cast + | AGGREGATION expression_math ")" [window_form] -> sql_aggregation + | "RANK"i "(" ")" window_form -> rank_expression + | "DENSE_RANK"i "(" ")" window_form -> dense_rank_expression + | "COALESCE"i "(" [(expression_math ",")*] expression_math ")" -> coalesce_expression + +window_form: "OVER"i "(" ["PARTITION"i "BY"i (partition_by ",")* partition_by] ["ORDER"i "BY"i (order ",")* order [ row_range_clause ] ] ")" + +partition_by: expression_math + +row_range_clause: ( ROWS | RANGE ) frame_extent +frame_extent: frame_between | frame_preceding +frame_between: "BETWEEN"i frame_bound "AND"i frame_bound +frame_bound: frame_preceding | frame_following | "CURRENT"i "ROW"i +frame_preceding: UNBOUNDED PRECEDING | integer_ PRECEDING +frame_following: UNBOUNDED FOLLOWING | integer_ FOLLOWING +RANGE: "RANGE"i +ROWS: "ROWS"i +UNBOUNDED: "UNBOUNDED"i +PRECEDING: "PRECEDING"i +FOLLOWING: "FOLLOWING"i + +when_then: "WHEN"i bool_expression "THEN"i expression_math +order: expression_math ["ASC"i] -> order_asc + | expression_math "DESC"i -> order_desc + +column_name: [name "."] name +?expression_product: expression_parens + | expression_product "*" expression_parens -> expression_mul + | expression_product "/" expression_parens -> expression_div + +?expression_parens: expression + | "(" expression_parens "*" expression ")" -> expression_mul + | "(" expression_parens "/" expression ")" -> expression_div + | "(" expression_parens "+" expression ")" -> expression_add + | "(" expression_parens "-" expression ")" -> expression_sub + +?expression: [name "."] (name | STAR) -> column_name + | literal + + +SELECT_CONSTRAINT.9: "ALL"i | "DISTINCT"i +TYPENAME: "object"i + | "varchar"i + | "integer"i + | "int16"i + | "smallint"i + | "int32"i + | "int64"i + | "int"i + | "bigint"i + | "float16"i + | "float32"i + | "float64"i + | "float"i + | "bool"i + | "datetime64"i + | "timestamp"i + | "time"i + | "date"i + | "category"i + | "string"i +AGGREGATION.8: ("sum("i | "avg("i | "min("i | "max("i | "count("i "distinct"i | "count("i) +alias: name -> alias_string +_window_name: name +limit_count: integer_ -> limit_count +skip_rows: integer_ +bool_expression: bool_parentheses + | bool_expression "AND"i bool_parentheses -> bool_and + | bool_expression "OR"i bool_parentheses -> bool_or +bool_parentheses: comparison_type + | "(" bool_expression "AND"i comparison_type ")" -> bool_and + | "(" bool_expression "OR"i comparison_type ")" -> bool_or +comparison_type: equals | not_equals | greater_than | less_than | greater_than_or_equal +| less_than_or_equal | between | in_expr | not_in_expr | subquery_in | is_null | is_not_null +equals: expression_math "=" expression_math +is_null: expression_math "is"i "null"i +is_not_null: expression_math "is"i "not"i "null"i +not_equals: expression_math ("<>" | "!=") expression_math +greater_than: expression_math ">" expression_math +less_than: expression_math "<" expression_math +greater_than_or_equal: expression_math ">=" expression_math +less_than_or_equal: expression_math "<=" expression_math +between: expression_math "BETWEEN"i expression_math "AND"i expression_math +in_expr: expression_math "IN"i "(" [expression_math ","]* expression_math ")" +subquery_in: expression_math "IN"i subquery +not_in_expr: expression_math "NOT"i "IN"i "(" [expression_math ","]* expression_math ")" +?literal: boolean -> bool + | number_expr -> number + | /'([^']|\s)+'|''/ -> string + | timestamp_expression -> timestamp_expression +boolean: "true"i -> true + | "false"i -> false +?number_expr: product + +?product: NUMBER + +integer_: /[1-9][0-9]*/ +STAR: "*" +window_definition: +timestamp_expression: "NOW"i "(" ")" -> datetime_now + | "TODAY"i "(" ")" -> date_today + | "TIMESTAMP"i "(" "'" date "'" "," "'" time "'" ")" -> custom_timestamp + +date: YEAR "-" MONTH "-" DAY +YEAR: /[0-9]{4}/ +MONTH: /[0-9]{2}/ +DAY: /[0-9]{2}/ +time: HOURS ":" MINUTES ":" SECONDS +HOURS: /[0-9]{2}/ +MINUTES: /[0-9]{2}/ +SECONDS: /[0-9]{2}/ +name: CNAME | ESCAPED_STRING + + + +%import common.ESCAPED_STRING +%import common.CNAME +%import common.NUMBER +%import common.WS +%import common.SQL_COMMENT +%import common.WS_INLINE + +%ignore WS +%ignore SQL_COMMENT diff --git a/outlines/models/mlxlm.py b/outlines/models/mlxlm.py index f561f269d2..57aa6f596f 100644 --- a/outlines/models/mlxlm.py +++ b/outlines/models/mlxlm.py @@ -9,7 +9,7 @@ from transformers import PreTrainedTokenizer from outlines.generate.api import GenerationParameters, SamplingParameters - from outlines.processors import BaseLogitsProcessor + from outlines.processors import OutlinesLogitsProcessor class MLXLM: @@ -120,7 +120,7 @@ def generate_step( temp: Optional[float], top_p: Optional[float], sampler: str, - logits_processor: "BaseLogitsProcessor", + logits_processor: "OutlinesLogitsProcessor", ) -> Generator[Tuple[int, float], None, None]: """ Adapted from @@ -135,7 +135,7 @@ def generate_step( top_p (float, optional): Nulceus sampling, higher means model considers more less likely words. sampler (str): The sampler string defined by SequenceGeneratorAdapter - logits_processor (BaseLogitsProcessor): Augment logits before sampling. + logits_processor (OutlinesLogitsProcessor): Augment logits before sampling. """ import mlx.core as mx import mlx_lm diff --git a/outlines/processors/__init__.py b/outlines/processors/__init__.py index 5c6a697ed6..22c10d9059 100644 --- a/outlines/processors/__init__.py +++ b/outlines/processors/__init__.py @@ -1,7 +1,7 @@ from .structured import ( - BaseLogitsProcessor, CFGLogitsProcessor, FSMLogitsProcessor, JSONLogitsProcessor, + OutlinesLogitsProcessor, RegexLogitsProcessor, ) diff --git a/outlines/processors/base_logits_processor.py b/outlines/processors/base_logits_processor.py index dabfd91b00..f829844a68 100644 --- a/outlines/processors/base_logits_processor.py +++ b/outlines/processors/base_logits_processor.py @@ -1,23 +1,30 @@ from abc import abstractmethod -from typing import List, Protocol, Union +from typing import TYPE_CHECKING, List, Protocol, Type, Union import numpy as np import torch from numpy.typing import NDArray +if TYPE_CHECKING: + import mlx.core as mx -def is_mlx_array(logits): + +Array = Union[NDArray, torch.Tensor, List, "mx.array"] + + +def is_mlx_array_type(array_type): try: import mlx.core as mx except ImportError: return False - return isinstance(logits, mx.array) + return issubclass(array_type, mx.array) -class BaseLogitsProcessor(Protocol): +class OutlinesLogitsProcessor(Protocol): """ Base class for logits processors which normalizes types of logits: - ndarray (used by llama-cpp-python), converted to torch.Tensor + - mlx.core.array (used by mlx-lm), converted to torch.Tensor - torch.Tensor (used by everything else) Normalization of types and conversion to torch.Tensor @@ -29,50 +36,100 @@ class BaseLogitsProcessor(Protocol): @abstractmethod def process_logits( - self, input_ids: List[int], logits: torch.Tensor + self, input_ids: List[List[int]], logits: torch.Tensor ) -> torch.Tensor: - ... + """ + input_ids and logits are always 2D tensors for handling a batch of sequences. + + - input_ids -> List[List[tokens]] + - logits.shape[0] -> 2D_Tensor[logits] + + Important to keep in mind when designing universal logits processors + - logits processors are only used once and never re-applied for a new sequence generator + - Some models only pass output_ids, some models such as llamacpp and transformers prefix with input_ids + - Some sampling methods, such as beam search, result in unstable sequence ordering in models like vLLM + """ + pass + @torch.no_grad() def __call__( self, - input_ids: Union[NDArray[np.int64], List[int], torch.Tensor], - logits: Union[NDArray[np.float32], torch.Tensor], - ) -> Union[NDArray[np.int64], torch.Tensor]: + input_ids: Array, + logits: Array, + ) -> Array: """ Apply logits processor - Unify type - - convert input_ids: either ndarray, List[int], or Tensor -> List[int] - - convert logits: either ndarray, mlx array, Tensor -> Tensor - Call process_logits() to perform business logic + + 1) Unify type + - convert input_ids: either ndarray, mlx array, List[int], or Tensor -> List[List[int]] + - convert logits: either ndarray, mlx array, or Tensor -> 2D float Tensor + 2) Unify shape, ensure logits and input_ids are 2D + 3) Call self.process_logits() to perform business logic + 4) Cast logits back to original array library type """ - with torch.no_grad(): - if not isinstance(input_ids, list): - input_ids = input_ids.tolist() - - if isinstance(logits, np.ndarray): - # Unify type, convert numpy array to Tensor - # from_numpy and .numpy() don't copy the data, it uses the same memory address - torch_logits = torch.from_numpy(logits) - processed_torch_logits = self.process_logits(input_ids, torch_logits) - return processed_torch_logits.detach().numpy() - - elif isinstance(logits, torch.Tensor): - return self.process_logits(input_ids, logits) - - elif is_mlx_array(logits): - # mlx -> torch -> mlx conversion docs: - # https://ml-explore.github.io/mlx/build/html/usage/numpy.html - import mlx.core as mx - - torch_logits = torch.from_dlpack(logits) - processed_torch_logits = self.process_logits(input_ids, torch_logits) - - # numpy doesn't support bfloat16, mlx doesn't support direct conversion from torch - logits_float32_numpy = processed_torch_logits.float().numpy() - return mx.array(logits_float32_numpy) - - else: - raise TypeError( - "LogitsProcessor must be called with either np.NDArray" - ", torch.Tensor, or mlx.core.array typed logits" - ) + + # ensure logits are torch Tensors + torch_logits = self._to_torch(logits) + + assert torch_logits.shape[:-1] == self._to_torch(input_ids).shape[:-1] + + # ensure input_ids are List + if not isinstance(input_ids, list): + input_ids = input_ids.tolist() # compatible with numpy, torch, and mlx + + # Guarantee passed as 2D Tensors, then covert back to original (1D or 2D) shape + if len(torch_logits.shape) == 2: + processed_logits = self.process_logits(input_ids, torch_logits) + elif len(torch_logits.shape) == 1: + processed_logits = self.process_logits( + [input_ids], torch_logits.unsqueeze(0) + ).squeeze(0) + + # return logits as passed array type + return self._from_torch(processed_logits, type(logits)) + + @staticmethod + def _to_torch(tensor_like: Array) -> torch.Tensor: + """Convert various types to torch.Tensor.""" + if isinstance(tensor_like, torch.Tensor): + return tensor_like + + elif isinstance(tensor_like, np.ndarray): + return torch.from_numpy(tensor_like) + + elif isinstance(tensor_like, list): + return torch.tensor(tensor_like) + + elif is_mlx_array_type(type(tensor_like)): + # mlx -> torch -> mlx conversion docs: + # https://ml-explore.github.io/mlx/build/html/usage/numpy.html + return torch.from_dlpack(tensor_like) + + else: + raise TypeError( + "LogitsProcessor must be called with either np.NDArray, " + "torch.Tensor, list, or mlx.core.array typed logits" + ) + + @staticmethod + def _from_torch(tensor: torch.Tensor, target_type: Type) -> Array: + """Convert torch.Tensor to the specified target type.""" + if target_type == torch.Tensor: + return tensor + + elif target_type == np.ndarray: + return tensor.detach().numpy() + + elif target_type == list: + return tensor.detach().tolist() + + elif is_mlx_array_type(target_type): + import mlx.core as mx + + # numpy doesn't support bfloat16, mlx doesn't support direct conversion from torch + return mx.array(tensor.float().numpy()) + + else: + raise TypeError( + f"Failed to convert torch tensors to target_type `{target_type}`" + ) diff --git a/outlines/processors/structured.py b/outlines/processors/structured.py index b8ef5b2da9..d037c679fc 100644 --- a/outlines/processors/structured.py +++ b/outlines/processors/structured.py @@ -24,24 +24,22 @@ limitations under the License. """ import math -from typing import TYPE_CHECKING, List, Optional, Type, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union -import numpy as np import torch -from numpy.typing import NDArray from pydantic import BaseModel from outlines.fsm.guide import CFGGuide, Guide, RegexGuide from outlines.fsm.json_schema import build_regex_from_schema from outlines.integrations.utils import convert_json_schema_to_str -from .base_logits_processor import BaseLogitsProcessor +from .base_logits_processor import OutlinesLogitsProcessor if TYPE_CHECKING: from outlines.models.tokenizer import Tokenizer -class FSMLogitsProcessor(BaseLogitsProcessor): +class FSMLogitsProcessor(OutlinesLogitsProcessor): """Bias generation using a finite state machine. Attributes @@ -63,13 +61,14 @@ def __init__(self, tokenizer: "Tokenizer", fsm: Guide): The finite state machine which is used to bias the logits. """ self.tokenizer = tokenizer - self._fsm_state = 0 + self._fsm_states: Dict[int, int] = {} self.fsm: Guide = fsm self._is_first_token = True + self._seq_start_idx: Optional[int] = None def process_logits( - self, input_ids: List[int], logits: torch.Tensor - ) -> NDArray[np.float32]: + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: """Use the FSM to bias the logits before sampling the next token. Parameters @@ -84,17 +83,31 @@ def process_logits( torch.Tensor The biased logits. """ + sequence_states: List[int] = [] # vector of states corresponding to `input_ids` + if self._is_first_token: self._is_first_token = False + self._seq_start_idx = len(input_ids[0]) + + self._fsm_states = {hash(tuple([])): 0} + sequence_states = [0] * len(input_ids) + else: - last_token = input_ids[-1] - self._fsm_state = self.fsm.get_next_state(self._fsm_state, last_token) + for seq_ids in input_ids: + prev_state_key = hash(tuple(seq_ids[self._seq_start_idx : -1])) + prev_state = self._fsm_states[prev_state_key] - allowed_tokens = self.fsm.get_next_instruction(self._fsm_state).tokens - allowed_tokens = torch.tensor(allowed_tokens, device=logits.device) + curr_state_key = hash(tuple(seq_ids[self._seq_start_idx :])) + curr_state = self.fsm.get_next_state(prev_state, seq_ids[-1]) + + self._fsm_states[curr_state_key] = curr_state + sequence_states.append(curr_state) mask = torch.full_like(logits, -math.inf) - mask[allowed_tokens] = logits[allowed_tokens] + for i, fsm_state in enumerate(sequence_states): + allowed_tokens = self.fsm.get_next_instruction(fsm_state).tokens + mask[i, allowed_tokens] = logits[i, allowed_tokens] + return mask def copy(self) -> "FSMLogitsProcessor": diff --git a/tests/cfg_samples/arithmetic/lots_of_ops.arithmetic.test b/tests/cfg_samples/arithmetic/lots_of_ops.arithmetic.test new file mode 100644 index 0000000000..dc65c21379 --- /dev/null +++ b/tests/cfg_samples/arithmetic/lots_of_ops.arithmetic.test @@ -0,0 +1 @@ +5+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-7-3-3-3-3-3-3-3-3-3-3-3-3-3-3-3-3-3-3-3-3+2+2+2+2+2+2+2+2-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-1-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2-2+7+7+7+7+7+7 diff --git a/tests/cfg_samples/arithmetic/simple_math.arithmetic.test b/tests/cfg_samples/arithmetic/simple_math.arithmetic.test new file mode 100644 index 0000000000..882f05c8da --- /dev/null +++ b/tests/cfg_samples/arithmetic/simple_math.arithmetic.test @@ -0,0 +1 @@ +(1 * 2) - (0.1 * 2 * 9.42) diff --git a/tests/cfg_samples/json/outlines.generate.samplers.mypy.json.test b/tests/cfg_samples/json/outlines.generate.samplers.mypy.json.test new file mode 100644 index 0000000000..1a328a9b6a --- /dev/null +++ b/tests/cfg_samples/json/outlines.generate.samplers.mypy.json.test @@ -0,0 +1,372 @@ +{ + ".class": "MypyFile", + "_fullname": "outlines.generate.samplers", + "future_import_flags": [], + "is_partial_stub_package": false, + "is_stub": false, + "names": { + ".class": "SymbolTable", + "Protocol": { + ".class": "SymbolTableNode", + "cross_ref": "typing.Protocol", + "kind": "Gdef" + }, + "Sampler": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "TypeInfo", + "_promote": [], + "abstract_attributes": [ + [ + "__call__", + 2 + ] + ], + "alt_promote": null, + "bases": [ + "builtins.object" + ], + "dataclass_transform_spec": null, + "declared_metaclass": null, + "defn": { + ".class": "ClassDef", + "fullname": "outlines.generate.samplers.Sampler", + "name": "Sampler", + "type_vars": [] + }, + "deletable_attributes": [], + "flags": [ + "is_abstract", + "is_protocol" + ], + "fullname": "outlines.generate.samplers.Sampler", + "has_param_spec_type": false, + "metaclass_type": "abc.ABCMeta", + "metadata": {}, + "module_name": "outlines.generate.samplers", + "mro": [ + "outlines.generate.samplers.Sampler", + "builtins.object" + ], + "names": { + ".class": "SymbolTable", + "__call__": { + ".class": "SymbolTableNode", + "kind": "Mdef", + "node": { + ".class": "FuncDef", + "abstract_status": 2, + "arg_kinds": [ + 0, + 0, + 0, + 0 + ], + "arg_names": [ + "self", + "logits", + "samples", + "rng" + ], + "dataclass_transform_spec": null, + "flags": [ + "is_trivial_body" + ], + "fullname": "outlines.generate.samplers.Sampler.__call__", + "name": "__call__", + "type": { + ".class": "CallableType", + "arg_kinds": [ + 0, + 0, + 0, + 0 + ], + "arg_names": [ + "self", + "logits", + "samples", + "rng" + ], + "arg_types": [ + "outlines.generate.samplers.Sampler", + { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + }, + "builtins.int", + { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + } + ], + "bound_args": [], + "def_extras": { + "first_arg": "self" + }, + "fallback": "builtins.function", + "from_concatenate": false, + "implicit": false, + "is_ellipsis_args": false, + "name": "__call__ of Sampler", + "ret_type": { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + }, + "type_guard": null, + "unpack_kwargs": false, + "variables": [] + } + } + } + }, + "self_type": null, + "slots": null, + "tuple_type": null, + "type_vars": [], + "typeddict_type": null + } + }, + "__annotations__": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "Var", + "flags": [ + "is_ready" + ], + "fullname": "outlines.generate.samplers.__annotations__", + "name": "__annotations__", + "type": { + ".class": "Instance", + "args": [ + "builtins.str", + { + ".class": "AnyType", + "missing_import_name": null, + "source_any": null, + "type_of_any": 6 + } + ], + "type_ref": "builtins.dict" + } + } + }, + "__doc__": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "Var", + "flags": [ + "is_ready" + ], + "fullname": "outlines.generate.samplers.__doc__", + "name": "__doc__", + "type": "builtins.str" + } + }, + "__file__": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "Var", + "flags": [ + "is_ready" + ], + "fullname": "outlines.generate.samplers.__file__", + "name": "__file__", + "type": "builtins.str" + } + }, + "__name__": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "Var", + "flags": [ + "is_ready" + ], + "fullname": "outlines.generate.samplers.__name__", + "name": "__name__", + "type": "builtins.str" + } + }, + "__package__": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "Var", + "flags": [ + "is_ready" + ], + "fullname": "outlines.generate.samplers.__package__", + "name": "__package__", + "type": "builtins.str" + } + }, + "greedy": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "FuncDef", + "abstract_status": 0, + "arg_kinds": [ + 0, + 0, + 2 + ], + "arg_names": [ + "logits", + "samples", + "_" + ], + "dataclass_transform_spec": null, + "flags": [], + "fullname": "outlines.generate.samplers.greedy", + "name": "greedy", + "type": { + ".class": "CallableType", + "arg_kinds": [ + 0, + 0, + 2 + ], + "arg_names": [ + "logits", + "samples", + "_" + ], + "arg_types": [ + { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + }, + "builtins.int", + { + ".class": "AnyType", + "missing_import_name": null, + "source_any": null, + "type_of_any": 1 + } + ], + "bound_args": [], + "def_extras": { + "first_arg": null + }, + "fallback": "builtins.function", + "from_concatenate": false, + "implicit": false, + "is_ellipsis_args": false, + "name": "greedy", + "ret_type": { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + }, + "type_guard": null, + "unpack_kwargs": false, + "variables": [] + } + } + }, + "multinomial": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "FuncDef", + "abstract_status": 0, + "arg_kinds": [ + 0, + 0, + 0 + ], + "arg_names": [ + "logits", + "samples", + "rng" + ], + "dataclass_transform_spec": null, + "flags": [], + "fullname": "outlines.generate.samplers.multinomial", + "name": "multinomial", + "type": { + ".class": "CallableType", + "arg_kinds": [ + 0, + 0, + 0 + ], + "arg_names": [ + "logits", + "samples", + "rng" + ], + "arg_types": [ + { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + }, + "builtins.int", + { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + } + ], + "bound_args": [], + "def_extras": { + "first_arg": null + }, + "fallback": "builtins.function", + "from_concatenate": false, + "implicit": false, + "is_ellipsis_args": false, + "name": "multinomial", + "ret_type": { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + }, + "type_guard": null, + "unpack_kwargs": false, + "variables": [] + } + } + }, + "torch": { + ".class": "SymbolTableNode", + "kind": "Gdef", + "node": { + ".class": "Var", + "flags": [ + "is_suppressed_import", + "is_ready", + "is_inferred" + ], + "fullname": "outlines.generate.samplers.torch", + "name": "torch", + "type": { + ".class": "AnyType", + "missing_import_name": "outlines.generate.samplers.torch", + "source_any": null, + "type_of_any": 3 + } + } + } + }, + "path": "/home/andrew/p/outlines/outlines/generate/samplers.py" +} diff --git a/tests/cfg_samples/json/simple_fruit.json.test b/tests/cfg_samples/json/simple_fruit.json.test new file mode 100644 index 0000000000..e8a4436250 --- /dev/null +++ b/tests/cfg_samples/json/simple_fruit.json.test @@ -0,0 +1,20 @@ +[ + { + "ID": "1", + "Name": "Andrew \"The Escaper\" Lapp", + "Age": "30", + "FavFruit": "Banana" + }, + { + "ID": "2", + "Name": "Mohammad", + "Age": "40", + "FavFruit": "\"Any Fruit As Long as It's In Quotes!\"" + }, + { + "ID": "3", + "Name": "Alice", + "Age": "61", + "FavFruit": "Peach" + } +] diff --git a/tests/cfg_samples/json/simple_fruit_no_indent.json.test b/tests/cfg_samples/json/simple_fruit_no_indent.json.test new file mode 100644 index 0000000000..9b7d319da2 --- /dev/null +++ b/tests/cfg_samples/json/simple_fruit_no_indent.json.test @@ -0,0 +1 @@ +[{"ID": "1", "Name": "Andrew", "Age": "30", "FavFruit": "Banana"}, {"ID": "2", "Name": "Mohammad", "Age": "40", "FavFruit": "Apple"}, {"ID": "3", "Name": "Alice", "Age": "61", "FavFruit": "Peach"}] diff --git a/tests/cfg_samples/sql_select/select_coalesce.sql.test b/tests/cfg_samples/sql_select/select_coalesce.sql.test new file mode 100644 index 0000000000..70c68f0d69 --- /dev/null +++ b/tests/cfg_samples/sql_select/select_coalesce.sql.test @@ -0,0 +1,2 @@ +-- SQL statement generated by ChatGPT +SELECT a.name, CASE WHEN a.age > 50 THEN 'Senior' ELSE 'Junior' END AS category, COALESCE(b.bonus, 0) FROM employees a LEFT JOIN bonuses b ON a.id = b.employee_id WHERE a.age BETWEEN 30 AND 60 diff --git a/tests/cfg_samples/sql_select/select_having.sql.test b/tests/cfg_samples/sql_select/select_having.sql.test new file mode 100644 index 0000000000..4fc8cf877c --- /dev/null +++ b/tests/cfg_samples/sql_select/select_having.sql.test @@ -0,0 +1,2 @@ +-- SQL statement generated by ChatGPT +SELECT a.department_id, COUNT(*) AS num_employees, AVG(b.salary) FROM employees a INNER JOIN salaries b ON a.id = b.employee_id GROUP BY a.department_id HAVING AVG(b.salary) > 50000 diff --git a/tests/cfg_samples/sql_select/select_many_features.sql.test b/tests/cfg_samples/sql_select/select_many_features.sql.test new file mode 100644 index 0000000000..a4164db13a --- /dev/null +++ b/tests/cfg_samples/sql_select/select_many_features.sql.test @@ -0,0 +1,29 @@ +-- complex SQL statement generated by ChatGPT +SELECT + T1."Employee Name", + T1.Department, + T2.TotalSales, + AVG(T3.Salary) OVER (PARTITION BY T1.Department ORDER BY T1."Employee Name" ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS AvgDeptSalary, + SUM(CASE WHEN T4.Month = 'January' THEN T4.Sales ELSE 0 END) AS JanuarySales, + RANK() OVER (ORDER BY T2.TotalSales DESC) AS SalesRank, + DENSE_RANK() OVER (PARTITION BY T1.Department ORDER BY T1."Employee Name") AS DeptRank, + COALESCE(T5.Bonus, 0) AS Bonus +FROM + Employees T1 +INNER JOIN + (SELECT EmployeeID, SUM(Sales) AS TotalSales FROM Sales GROUP BY EmployeeID) T2 ON T1.EmployeeID = T2.EmployeeID +LEFT JOIN --LEFT JOIN -- TODO: Fix LALR so LEFT JOIN is legal + Salaries T3 ON T1.EmployeeID = T3.EmployeeID AND T3.Year = 2023 +LEFT OUTER JOIN + Sales T4 ON T1.EmployeeID = T4.EmployeeID +LEFT OUTER JOIN + (SELECT EmployeeID, SUM(Bonus) AS Bonus FROM Bonuses WHERE Year = 2023 GROUP BY EmployeeID) T5 ON T1.EmployeeID = T5.EmployeeID +WHERE + T1.Department IS NOT NULL AND T2.TotalSales > 10000 +GROUP BY + T1."Employee Name", T1.Department, T2.TotalSales, T5.Bonus + HAVING + AVG(T3.Salary) > 50000 +ORDER BY + SalesRank, T1.Department, T1."Employee Name" +LIMIT 10 diff --git a/tests/cfg_samples/sql_select/select_minimal_lalr1.sql.test b/tests/cfg_samples/sql_select/select_minimal_lalr1.sql.test new file mode 100644 index 0000000000..8a275802fa --- /dev/null +++ b/tests/cfg_samples/sql_select/select_minimal_lalr1.sql.test @@ -0,0 +1,5 @@ +SELECT * +FROM foo +HAVING +AVG(foo.bar) > 1 +ORDER BY foo.baz diff --git a/tests/cfg_samples/sql_select/select_nested_subquery.sql.test b/tests/cfg_samples/sql_select/select_nested_subquery.sql.test new file mode 100644 index 0000000000..81d4a2b064 --- /dev/null +++ b/tests/cfg_samples/sql_select/select_nested_subquery.sql.test @@ -0,0 +1,21 @@ +-- SQL statement generated by ChatGPT +SELECT + a.name, + avg_salaries.avg_salary +FROM + departments a +LEFT JOIN ( + SELECT + b.department_id, + AVG(b.salary) AS avg_salary + FROM + employees b + GROUP BY + b.department_id +) avg_salaries ON a.department_id = avg_salaries.department_id +WHERE + a.location = 'New York' +ORDER BY + a.name, + avg_salaries.avg_salary DESC +LIMIT 5 OFFSET 10 diff --git a/tests/cfg_samples/sql_select/select_order.sql.test b/tests/cfg_samples/sql_select/select_order.sql.test new file mode 100644 index 0000000000..268340dbee --- /dev/null +++ b/tests/cfg_samples/sql_select/select_order.sql.test @@ -0,0 +1,3 @@ +SELECT * +FROM myTable +WHERE myTable.foo = 'bar' diff --git a/tests/cfg_samples/sql_select/select_simple.sql.test b/tests/cfg_samples/sql_select/select_simple.sql.test new file mode 100644 index 0000000000..eee760e795 --- /dev/null +++ b/tests/cfg_samples/sql_select/select_simple.sql.test @@ -0,0 +1 @@ +SELECT * FROM foo diff --git a/tests/cfg_samples/sql_select/select_union.sql.test b/tests/cfg_samples/sql_select/select_union.sql.test new file mode 100644 index 0000000000..f7913afb9c --- /dev/null +++ b/tests/cfg_samples/sql_select/select_union.sql.test @@ -0,0 +1,5 @@ +-- SQL statement generated by ChatGPT +SELECT a.department, COUNT(*) FROM employees a GROUP BY a.department +UNION ALL +SELECT b.department, COUNT(*) FROM managers b GROUP BY b.department +ORDER BY 1 diff --git a/tests/fsm/test_integration_cfg.py b/tests/fsm/test_integration_cfg.py new file mode 100644 index 0000000000..234f90ab34 --- /dev/null +++ b/tests/fsm/test_integration_cfg.py @@ -0,0 +1,72 @@ +from outlines.fsm.guide import CFGGuide, Generate +import outlines.grammars as grammars +import outlines.models as models + +from pathlib import Path + +import pytest + + +@pytest.fixture(scope="session") +def tokenizer_bpe(): + return models.transformers( + "hf-internal-testing/tiny-random-gpt2", device="cpu" + ).tokenizer + + +@pytest.fixture(scope="session") +def tokenizer_character_level(): + return models.transformers("google/byt5-small", device="cpu").tokenizer + + +TOKENIZERS = ["tokenizer_bpe"] # , "tokenizer_character_level"] + + +# Collects all samples within cfg_samples/ and makes adding +# a test case as easy as adding a valid sample to cfg_samples/ +all_samples = {} +examples_path = Path(__file__).parent.parent / "cfg_samples" +for sample_collection_path in examples_path.iterdir(): + grammar_name = sample_collection_path.name + grammar = getattr(grammars, grammar_name) + for sample_path in sample_collection_path.iterdir(): + test_name = f"{grammar_name}_{sample_path.name}" + with open(sample_path) as f: + all_samples[test_name] = (grammar, f.read().rstrip("\n")) + + +@pytest.mark.parametrize("sample_name", all_samples.keys()) +@pytest.mark.parametrize("tokenizer_name", TOKENIZERS) +def test_cfg_sample_valid(request, sample_name, tokenizer_name): + tokenizer = request.getfixturevalue(tokenizer_name) + + cfg, sample = all_samples[sample_name] + cfg_guide = CFGGuide(cfg, tokenizer) + + # TODO: assert that the sample is valid under the grammar using + # pure lark, if its not raise an appropriate exception + + sample_token_ids = tokenizer.encode(sample)[0][0] + assert len(sample_token_ids.shape) == 1 # ensure we're encoding in the desired shape for this test + + state = 0 + for i, token_id in enumerate(sample_token_ids): + next_instruction = cfg_guide.get_next_instruction(state) + if token_id not in next_instruction.tokens: + processed_str = tokenizer.decode([sample_token_ids[:i]])[0] + remaining_str = tokenizer.decode([sample_token_ids[i:]])[0] + if next_instruction.tokens == [tokenizer.eos_token_id]: + error_label = "CFGGuide required EOS early" + else: + expected = tokenizer.decode(next_instruction.tokens) + error_label = f"Mismatched expectations, Guide expected {expected}" + raise Exception( + f"{error_label}\n" + f"processed:\n```{processed_str}```\n" + f"remaining:\n```{remaining_str}```" + ) + next_instruction.tokens + state = cfg_guide.get_next_state(state, token_id) + + final_instruction = cfg_guide.get_next_instruction(state) + assert tokenizer.eos_token_id in final_instruction diff --git a/tests/generate/test_generate.py b/tests/generate/test_generate.py index 1f1a3aea27..111f8f93db 100644 --- a/tests/generate/test_generate.py +++ b/tests/generate/test_generate.py @@ -1,9 +1,11 @@ +import contextlib import re import pytest import outlines.generate as generate import outlines.models as models +import outlines.samplers as samplers @pytest.fixture(scope="session") @@ -20,35 +22,160 @@ def model_mlxlm(tmp_path_factory): @pytest.fixture(scope="session") -def model_transformers(tmp_path_factory): - return models.transformers("Locutusque/TinyMistral-248M-v2-Instruct", device="cpu") +def model_transformers_random(tmp_path_factory): + return models.transformers("hf-internal-testing/tiny-random-gpt2", device="cpu") -@pytest.mark.parametrize( - "model_fixture", - ("model_llamacpp", "model_mlxlm", "model_transformers"), +@pytest.fixture(scope="session") +def model_transformers_opt125m(tmp_path_factory): + return models.transformers("facebook/opt-125m", device="cpu") + + +ALL_MODEL_FIXTURES = ( + "model_llamacpp", + "model_mlxlm", + "model_transformers_random", + "model_transformers_opt125m", ) -def test_generate_text(request, model_fixture): + + +NOT_IMPLEMENTED = { + "batch": ["model_llamacpp"], + "stream": ["model_vllm"], + "beam_search": ["model_llamacpp"], + "multiple_samples": ["model_llamacpp"], +} + + +def enforce_not_implemented(model_fixture, *task_names): + """ + Per `NOT_IMPLEMENTED`, mapping, if a model hasn't implemented a task, + assert an NotImplementedError is raised. Otherwise, run normally + """ + for task_name in task_names: + if model_fixture in NOT_IMPLEMENTED.get(task_name, []): + return pytest.raises(NotImplementedError) + else: + return contextlib.nullcontext() + + +REGEX_PATTERNS = [ + "(123456789)|(abcdefghijklmnop)", + "abc*", + "\\+?[1-9][0-9]{7,14}", + r"([a-z]{10})@([a-z]{5})\.([a-z]{3})", +] + + +@pytest.mark.parametrize("sampler_name", ("greedy", "multinomial", "beam_search")) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_text(request, model_fixture, sampler_name): + model = request.getfixturevalue(model_fixture) + generator = generate.text(model, getattr(samplers, sampler_name)()) + with enforce_not_implemented(model_fixture, sampler_name): + res = generator("test", max_tokens=10) + assert isinstance(res, str) + + +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_batch_text(request, model_fixture): model = request.getfixturevalue(model_fixture) generator = generate.text(model) - res = generator("test", max_tokens=10) - assert isinstance(res, str) + with enforce_not_implemented(model_fixture, "batch"): + res = generator(["test", "test2"], max_tokens=10) + assert isinstance(res, list) + assert isinstance(res[0], str) -@pytest.mark.parametrize( - "model_fixture", - ("model_llamacpp", "model_mlxlm", "model_transformers"), -) -@pytest.mark.parametrize( - "pattern", - ( - "[0-9]", - "abc*", - "\\+?[1-9][0-9]{7,14}", - ), -) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_text_stream(request, model_fixture): + model = request.getfixturevalue(model_fixture) + generator = generate.text(model) + with enforce_not_implemented(model_fixture, "stream"): + for token in generator.stream("a b c ", max_tokens=10): + assert isinstance(token, str) + + +@pytest.mark.parametrize("pattern", REGEX_PATTERNS) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) def test_generate_regex(request, model_fixture, pattern): model = request.getfixturevalue(model_fixture) generator = generate.regex(model, pattern) res = generator("foobarbaz", max_tokens=20) - assert re.match(pattern, res) is not None, res + assert re.fullmatch(pattern, res) is not None, res + + +@pytest.mark.parametrize("pattern", REGEX_PATTERNS) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_regex_stream(request, model_fixture, pattern): + model = request.getfixturevalue(model_fixture) + generator = generate.regex(model, pattern) + with enforce_not_implemented(model_fixture, "stream"): + output = "" + for token in generator.stream("output:", max_tokens=20): + output += token + assert re.fullmatch(pattern, output) is not None, output + + +@pytest.mark.parametrize("pattern", REGEX_PATTERNS) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_regex_batch_stream(request, model_fixture, pattern): + model = request.getfixturevalue(model_fixture) + generator = generate.regex(model, pattern) + with enforce_not_implemented(model_fixture, "batch", "stream"): + outputs = ["", ""] + for tokens in generator.stream(["input 0", "input 1"], max_tokens=20): + outputs[0] += tokens[0] + outputs[1] += tokens[1] + for output in outputs: + assert re.fullmatch(pattern, output) is not None, output + + +@pytest.mark.parametrize("pattern", REGEX_PATTERNS) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_regex_batch(request, model_fixture, pattern): + """Ensure batch requests work and fsm order is maintained""" + model = request.getfixturevalue(model_fixture) + generator = generate.regex(model, pattern) + with enforce_not_implemented(model_fixture, "batch"): + outputs = generator(["abc", "123", "123bce", "33aa"], max_tokens=20) + for output in outputs: + assert re.fullmatch(pattern, output) is not None, output + + +@pytest.mark.parametrize("pattern", REGEX_PATTERNS) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_regex_single_multinomial(request, model_fixture, pattern): + """Ensure batch requests work and fsm order is maintained""" + model = request.getfixturevalue(model_fixture) + generator = generate.regex(model, pattern, sampler=samplers.multinomial(4)) + with enforce_not_implemented(model_fixture, "multiple_samples"): + output_sample_groups = generator("single input", max_tokens=40) + for output in output_sample_groups: + assert re.fullmatch(pattern, output) is not None, output + + +@pytest.mark.parametrize("pattern", REGEX_PATTERNS) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_regex_batch_multinomial(request, model_fixture, pattern): + """Ensure batch requests work and fsm order is maintained""" + model = request.getfixturevalue(model_fixture) + generator = generate.regex(model, pattern, sampler=samplers.multinomial(4)) + with enforce_not_implemented(model_fixture, "batch", "multiple_samples"): + output_batch_groups = generator(["abc", "123", "123bce", "33aa"], max_tokens=40) + for output_sample_groups in output_batch_groups: + for output in output_sample_groups: + assert re.fullmatch(pattern, output) is not None, output + + +@pytest.mark.parametrize("pattern", REGEX_PATTERNS) +@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES) +def test_generate_regex_batch_beam_search(request, model_fixture, pattern): + """Ensure batch requests work and fsm order is maintained""" + model = request.getfixturevalue(model_fixture) + generator = generate.regex(model, pattern, sampler=samplers.beam_search(4)) + with enforce_not_implemented(model_fixture, "batch", "multiple_samples"): + output_batch_groups = generator(["abc", "123", "123bce", "33aa"], max_tokens=40) + for output_sample_groups in output_batch_groups: + for output in output_sample_groups: + assert re.fullmatch(pattern, output) is not None, output