Spaces:
Sleeping
Sleeping
| from copy import copy | |
| from functools import partial | |
| from outlines.fsm.guide import RegexGuide | |
| from pydantic import BaseModel | |
| from transformers import PreTrainedTokenizerBase | |
| def merge_successive_transitions(states_to_token_maps: dict[int, dict[int, int]], i, j): | |
| states_to_token_maps = dict(states_to_token_maps) | |
| transitions_i = {(s1, states_to_token_maps[s1][i]) for s1 in states_to_token_maps if i in states_to_token_maps[s1]} | |
| transitions_j = {(s1, states_to_token_maps[s1][j]) for s1 in states_to_token_maps if j in states_to_token_maps[s1]} | |
| transitions_i, transitions_j = dict(transitions_i - transitions_j), dict(transitions_j - transitions_i) | |
| for s1, s2 in transitions_i.items(): | |
| while s2 in transitions_j: | |
| s2 = transitions_j[s2] | |
| if s2 != transitions_i[s1]: | |
| states_to_token_maps[s1] = dict(states_to_token_maps[s1]) | |
| states_to_token_maps[s1][i] = s2 | |
| return states_to_token_maps | |
| def replace_transitions(states_to_token_maps: dict[int, dict[int, int]], i, j): | |
| states_to_token_maps = dict(states_to_token_maps) | |
| transitions_i = {(s1, states_to_token_maps[s1][i]) for s1 in states_to_token_maps if i in states_to_token_maps[s1]} | |
| transitions_j = {(s1, states_to_token_maps[s1][j]) for s1 in states_to_token_maps if j in states_to_token_maps[s1]} | |
| transitions_i, transitions_j = dict(transitions_i - transitions_j), dict(transitions_j - transitions_i) | |
| for s1, s2 in transitions_i.items(): | |
| if s2 != transitions_j.get(s1): | |
| states_to_token_maps[s1] = dict(states_to_token_maps[s1]) | |
| if s1 in transitions_j: | |
| states_to_token_maps[s1][i] = transitions_j[s1] | |
| else: | |
| states_to_token_maps[s1].pop(i) | |
| states_to_token_maps[s1][j] = s2 | |
| return states_to_token_maps | |
| def find_paths_with_transitions(states_to_token_maps: dict[int, dict[int, int]], transitions: list[int]) -> list[list[int]]: | |
| possible_s0 = {s0 for s0 in states_to_token_maps if transitions[0] in states_to_token_maps[s0]} | |
| possible_s1 = {s1 for s1 in states_to_token_maps if transitions[1] in states_to_token_maps[s1]} - possible_s0 | |
| starts = sorted( | |
| s0 for s0 in possible_s0 | |
| if states_to_token_maps[s0][transitions[0]] in possible_s1 | |
| ) | |
| paths = [[start] for start in starts] | |
| for path in paths: | |
| for i in transitions: | |
| if i in states_to_token_maps[path[-1]]: | |
| path.append(states_to_token_maps[path[-1]][i]) | |
| else: | |
| break | |
| return [path for path in paths if len(path) == len(transitions) + 1] | |
| def replace_fields(fsm: RegexGuide, model: BaseModel, new_fields: list[str], tokenizer: PreTrainedTokenizerBase, make_infinite_loop: bool = False) -> RegexGuide: | |
| assert len(new_fields) <= len(model.model_fields) | |
| sttm = dict(fsm.states_to_token_maps) | |
| encode = partial(tokenizer.encode, add_special_tokens=False) | |
| quote = encode('"')[0] | |
| # Let's replace the placeholder fields from the model in the finite state model by the new fields | |
| for orig_field, new_field in zip(model.model_fields, new_fields): | |
| orig_field_tokens = [encode(orig_field_char)[0] for orig_field_char in orig_field] | |
| new_field_tokens = encode(new_field) | |
| assert len(new_field_tokens) <= len(orig_field_tokens) | |
| # Merge transitions until we have number of transitions = number of tokens in the field name | |
| for k in reversed(range(len(new_field_tokens), len(orig_field_tokens))): | |
| sttm = merge_successive_transitions(sttm, orig_field_tokens[k - 1], orig_field_tokens[k]) | |
| # Replace the token ids in the transitions with the ones of the new field name | |
| for k in range(len(new_field_tokens)): | |
| sttm = replace_transitions(sttm, orig_field_tokens[k], new_field_tokens[k]) | |
| if len(new_fields) < len(model.model_fields) or make_infinite_loop: | |
| # Set the last field last state to generate less than the number of fields in the model | |
| # We need to do this for every possible path | |
| # e.g. multiple paths are used to count items when setting a min/max length | |
| orig_last_field = list(model.model_fields)[-1] | |
| new_last_field = new_fields[-1] | |
| orig_last_field_paths = find_paths_with_transitions(sttm, [quote] + [encode(c)[0] for c in orig_last_field]) | |
| new_last_field_paths = find_paths_with_transitions(sttm, [quote] + encode(new_last_field)) | |
| if make_infinite_loop: # this is a hack to loop on the same states over and over again | |
| orig_last_field_paths = [orig_last_field_paths[0]] * len(orig_last_field_paths) | |
| for orig_last_field_path, new_last_field_path in zip( | |
| orig_last_field_paths, | |
| new_last_field_paths | |
| ): | |
| orig_last_field_last_state = orig_last_field_path[-1] | |
| new_last_field_second_last_state = new_last_field_path[-2] | |
| sttm[new_last_field_second_last_state] = dict(sttm[new_last_field_second_last_state]) | |
| sttm[new_last_field_second_last_state][encode(new_last_field)[-1]] = orig_last_field_last_state | |
| fsm = copy(fsm) | |
| fsm.states_to_token_maps = sttm | |
| return fsm | |