zodiac.graph

  1#  # # <!-- // /*  SPDX-License-Identifier: MPL-2.0*/ -->
  2#  # # <!-- // /*  d a r k s h a p e s */ -->
  3
  4
  5import sys
  6import os
  7import networkx as nx
  8from typing import Optional
  9from nnll.monitor.file import dbug, dbuq
 10from zodiac.providers.pools import register_models  # leaving here for mocking
 11
 12sys.path.append(os.getcwd())
 13nfo = print
 14
 15
 16class IntentProcessor:
 17    intent_graph: Optional[dict[nx.Graph]] = None
 18    coord_path: Optional[list[str]] = None
 19    registry_entries: Optional[list[dict[dict]]] = None
 20    models: Optional[list[tuple[str]]] = None
 21    weight_idx: Optional[list[str]] = None
 22    # additional_model_names: dict = None
 23
 24    def __init__(self, intent_graph: nx.MultiDiGraph = nx.MultiDiGraph()) -> None:
 25        """
 26        Create instance of graph processor & initialize objectieves for tracing paths\n
 27        :param nx_graph:Preassembled graph of models to substitute, default uses nx.MultiDiGraph()
 28
 29        ========================================================\n
 30        ### GIVEN\n
 31        A : The list of `VALID CONVERSIONS` contains all of Zodiac's supported generative modalities\n
 32        B : The graph is populated directly from the contents of the list in A\n
 33        Thus: All possible node start and end points listed in A are included in graph B.\n
 34        Therefore : It is impossible to call a node that does not exist.\n
 35        """
 36        from zodiac.providers.constants import VALID_CONVERSIONS
 37
 38        self.intent_graph = intent_graph
 39        self.intent_graph.add_nodes_from(VALID_CONVERSIONS)
 40
 41    async def calc_graph(self, registry_entries: Optional[list] = None) -> None:
 42        """Generate graph of coordinate pairs from valid conversions\n
 43        Model libraries are auto-detected from cache loading\n
 44        :param registry_data: Registry function or method of calling registry, defaults to
 45        :return: Graph modeling all current ML/AI tasks appended with model data
 46
 47        ========================================================\n
 48        ### GIVEN\n
 49        A : The set of all models M on the executing system\n
 50        B : P is the randomly distributed set of start and end points required to graph M\n
 51        Thus: Because of the randomness of B, the set P is unlikely to construct a complete graph attached all available points.\n
 52        Therefore : While we can trust a node exists, we **CANNOT** trust the system has an edge to reach it\n
 53        """
 54        # import asyncio
 55
 56        if not registry_entries:
 57            registry_entries = await register_models()
 58        nfo("Building graph...")
 59
 60        if registry_entries is None:
 61            nfo("Registry error, graph attributes not applied.")
 62        elif len(self.intent_graph.edges) > 0:
 63            nfo("Edges already calculated")
 64            return self.intent_graph
 65        else:
 66            for model in registry_entries:
 67                try:
 68                    self.intent_graph.add_edges_from(model.available_tasks, entry=model, weight=1.0)
 69                except AttributeError as error_log:
 70                    dbug(error_log)
 71                    nfo("Error: Registry initialized but not populated with data. Graph could not create edges.")
 72
 73        nfo("Complete {self.intent_graph}")
 74        return self.intent_graph
 75
 76    def set_path(self, mode_in: str, mode_out: str) -> None:
 77        """Find a valid path from current state (mode_in) to designated state (mode_out)\n
 78        :param mode_in: Input prompt type or starting state/states
 79        :type mode_in: str
 80        :param mode_out: The user-selected ending-state
 81        :type mode_out: str
 82        """
 83
 84        if nx.has_path(self.intent_graph, mode_in, mode_out):  # Ensure path exists (otherwise 'bidirectional' may loop infinitely)
 85            # Self loops in the multidirected graph complete themselves
 86            # In practice, this means often the same model can be used to compute prompt input and response output
 87            # Unfortunately, this doesn't always work in all modalities, ex. Image to Image.
 88            # This condition is meant to solve the case of non-text self-loop edge being an incomplete transformation
 89
 90            if mode_in == mode_out and mode_in != "text":  # Its not a great solution, but it works for the moment
 91                orig_mode_out = mode_out
 92                mode_out = "text"
 93                self.coord_path = nx.bidirectional_shortest_path(self.intent_graph, mode_in, mode_out)
 94                self.coord_path.append(orig_mode_out)
 95            else:
 96                self.coord_path = nx.bidirectional_shortest_path(self.intent_graph, mode_in, mode_out)
 97                if len(self.coord_path) == 1:
 98                    self.coord_path.append(mode_out)  # this behaviour likely to change in future
 99
100        else:
101            nfo("No Path available...\n")
102
103    def set_registry_entries(self) -> None:
104        """Populate models list for text fields
105        Check if model has been adjusted, if so adjust list
106        1.0 weight bottom, <1.0 weight top"""
107
108        try:
109            self.registry_entries = self.pull_path_entries(self.intent_graph, self.coord_path)
110        except KeyError as error_log:
111            dbug(error_log)
112            return ["", ""]
113        idx = 0
114        self.models = []
115
116        if self.registry_entries:
117            for edge, registry in enumerate(self.registry_entries):
118                model = registry["entry"].model
119                dbuq(f"node {edge}")
120                adj_model = (os.path.basename(model), edge)
121                self.models.append(adj_model)
122            self.weight_idx = self.weight_idx or []
123            for model in self.weight_idx:
124                if model in self.models:
125                    self.models.remove(model)
126                    adj_model = (f"*{model[0]}", model[1])
127                    self.models.insert(idx, adj_model)
128                    idx += 1
129
130    def edit_weight(self, edge_number: str, mode_in: str, mode_out: str) -> None:
131        """Determine entry edge, determine index, then adjust weight\n
132        :param edge_number: Text pattern from `models` class attribute to identify the model by
133        :param mode_in: The conversion type, representing a source graph node
134        :param mode_out: The target type, , representing a source graph node
135        :raises ValueError: No models fit the request
136        """
137
138        self.weight_idx = self.weight_idx or []
139
140        try:
141            if not nx.has_path(self.intent_graph, mode_in, mode_out):
142                raise KeyError()
143            model = self.intent_graph[mode_in][mode_out][edge_number]["entry"].model
144        except KeyError as error_log:
145            nfo(
146                f"Failed to adjust weight of '{edge_number}' within registry contents \
147                '{self.intent_graph} {mode_in} {mode_out}'. Model or registry entry not found. "
148            )
149            dbug(error_log)
150            return self.set_registry_entries()
151
152        weight = self.intent_graph[mode_in][mode_out][edge_number]["weight"]
153        item = (os.path.basename(model), edge_number)
154        nfo(f" model : {model}  weight: {weight} ")
155
156        if weight < 1.0:
157            self.intent_graph[mode_in][mode_out][edge_number]["weight"] = round(weight + 0.1, 1)
158            self.models = [((f"*{os.path.basename(model)}", edge_number))]
159            if item in self.weight_idx:
160                self.weight_idx.remove(item)
161        else:
162            self.intent_graph[mode_in][mode_out][edge_number]["weight"] = round(weight - 0.1, 1)
163            self.weight_idx.append(item)
164        self.set_registry_entries()
165
166    def pull_path_entries(self, nx_graph: nx.Graph, traced_path: list[tuple]) -> None:
167        """Create operating instructions from user input
168        Trace the next hop along the path, collect all compatible models
169        Set current model based on weight and next available"""
170
171        registry_entries = []
172        if traced_path is not None and nx.has_path(nx_graph, traced_path[0], traced_path[1]):
173            registry_entries = [  # ruff : noqa
174                nx_graph[traced_path[index]][traced_path[index + 1]][hop]  #
175                for index in range(len(traced_path) - 1)  #
176                for hop in nx_graph[traced_path[index]][traced_path[index + 1]]  #
177            ]
178        return registry_entries