kumoai 2.13.0.dev202511261731__cp310-cp310-win_amd64.whl → 2.13.0.dev202512040252__cp310-cp310-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- kumoai/_version.py +1 -1
- kumoai/connector/utils.py +23 -2
- kumoai/experimental/rfm/__init__.py +20 -45
- kumoai/experimental/rfm/backend/__init__.py +0 -0
- kumoai/experimental/rfm/backend/local/__init__.py +38 -0
- kumoai/experimental/rfm/backend/local/table.py +109 -0
- kumoai/experimental/rfm/backend/snow/__init__.py +35 -0
- kumoai/experimental/rfm/backend/snow/table.py +115 -0
- kumoai/experimental/rfm/backend/sqlite/__init__.py +30 -0
- kumoai/experimental/rfm/backend/sqlite/table.py +101 -0
- kumoai/experimental/rfm/base/__init__.py +10 -0
- kumoai/experimental/rfm/base/column.py +66 -0
- kumoai/experimental/rfm/base/source.py +18 -0
- kumoai/experimental/rfm/{local_table.py → base/table.py} +134 -139
- kumoai/experimental/rfm/{local_graph.py → graph.py} +287 -62
- kumoai/experimental/rfm/infer/__init__.py +6 -0
- kumoai/experimental/rfm/infer/dtype.py +79 -0
- kumoai/experimental/rfm/infer/pkey.py +126 -0
- kumoai/experimental/rfm/infer/time_col.py +62 -0
- kumoai/experimental/rfm/local_graph_sampler.py +42 -1
- kumoai/experimental/rfm/local_graph_store.py +13 -27
- kumoai/experimental/rfm/rfm.py +6 -16
- kumoai/experimental/rfm/sagemaker.py +11 -3
- kumoai/kumolib.cp310-win_amd64.pyd +0 -0
- kumoai/testing/decorators.py +1 -1
- {kumoai-2.13.0.dev202511261731.dist-info → kumoai-2.13.0.dev202512040252.dist-info}/METADATA +9 -8
- {kumoai-2.13.0.dev202511261731.dist-info → kumoai-2.13.0.dev202512040252.dist-info}/RECORD +30 -18
- kumoai/experimental/rfm/utils.py +0 -344
- {kumoai-2.13.0.dev202511261731.dist-info → kumoai-2.13.0.dev202512040252.dist-info}/WHEEL +0 -0
- {kumoai-2.13.0.dev202511261731.dist-info → kumoai-2.13.0.dev202512040252.dist-info}/licenses/LICENSE +0 -0
- {kumoai-2.13.0.dev202511261731.dist-info → kumoai-2.13.0.dev202512040252.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
import re
|
|
2
|
+
import warnings
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import pandas as pd
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def infer_primary_key(
|
|
9
|
+
table_name: str,
|
|
10
|
+
df: pd.DataFrame,
|
|
11
|
+
candidates: list[str],
|
|
12
|
+
) -> Optional[str]:
|
|
13
|
+
r"""Auto-detect potential primary key column.
|
|
14
|
+
|
|
15
|
+
Args:
|
|
16
|
+
table_name: The table name.
|
|
17
|
+
df: The pandas DataFrame to analyze.
|
|
18
|
+
candidates: A list of potential candidates.
|
|
19
|
+
|
|
20
|
+
Returns:
|
|
21
|
+
The name of the detected primary key, or ``None`` if not found.
|
|
22
|
+
"""
|
|
23
|
+
# A list of (potentially modified) table names that are eligible to match
|
|
24
|
+
# with a primary key, i.e.:
|
|
25
|
+
# - UserInfo -> User
|
|
26
|
+
# - snakecase <-> camelcase
|
|
27
|
+
# - camelcase <-> snakecase
|
|
28
|
+
# - plural <-> singular (users -> user, eligibilities -> eligibility)
|
|
29
|
+
# - verb -> noun (qualifying -> qualify)
|
|
30
|
+
_table_names = {table_name}
|
|
31
|
+
if table_name.lower().endswith('_info'):
|
|
32
|
+
_table_names.add(table_name[:-5])
|
|
33
|
+
elif table_name.lower().endswith('info'):
|
|
34
|
+
_table_names.add(table_name[:-4])
|
|
35
|
+
|
|
36
|
+
table_names = set()
|
|
37
|
+
for _table_name in _table_names:
|
|
38
|
+
table_names.add(_table_name.lower())
|
|
39
|
+
snakecase = re.sub(r'(.)([A-Z][a-z]+)', r'\1_\2', _table_name)
|
|
40
|
+
snakecase = re.sub(r'([a-z0-9])([A-Z])', r'\1_\2', snakecase)
|
|
41
|
+
table_names.add(snakecase.lower())
|
|
42
|
+
camelcase = _table_name.replace('_', '')
|
|
43
|
+
table_names.add(camelcase.lower())
|
|
44
|
+
if _table_name.lower().endswith('s'):
|
|
45
|
+
table_names.add(_table_name.lower()[:-1])
|
|
46
|
+
table_names.add(snakecase.lower()[:-1])
|
|
47
|
+
table_names.add(camelcase.lower()[:-1])
|
|
48
|
+
else:
|
|
49
|
+
table_names.add(_table_name.lower() + 's')
|
|
50
|
+
table_names.add(snakecase.lower() + 's')
|
|
51
|
+
table_names.add(camelcase.lower() + 's')
|
|
52
|
+
if _table_name.lower().endswith('ies'):
|
|
53
|
+
table_names.add(_table_name.lower()[:-3] + 'y')
|
|
54
|
+
table_names.add(snakecase.lower()[:-3] + 'y')
|
|
55
|
+
table_names.add(camelcase.lower()[:-3] + 'y')
|
|
56
|
+
elif _table_name.lower().endswith('y'):
|
|
57
|
+
table_names.add(_table_name.lower()[:-1] + 'ies')
|
|
58
|
+
table_names.add(snakecase.lower()[:-1] + 'ies')
|
|
59
|
+
table_names.add(camelcase.lower()[:-1] + 'ies')
|
|
60
|
+
if _table_name.lower().endswith('ing'):
|
|
61
|
+
table_names.add(_table_name.lower()[:-3])
|
|
62
|
+
table_names.add(snakecase.lower()[:-3])
|
|
63
|
+
table_names.add(camelcase.lower()[:-3])
|
|
64
|
+
|
|
65
|
+
scores: list[tuple[str, int]] = []
|
|
66
|
+
for col_name in candidates:
|
|
67
|
+
col_name_lower = col_name.lower()
|
|
68
|
+
|
|
69
|
+
score = 0
|
|
70
|
+
|
|
71
|
+
if col_name_lower == 'id':
|
|
72
|
+
score += 4
|
|
73
|
+
|
|
74
|
+
for table_name_lower in table_names:
|
|
75
|
+
|
|
76
|
+
if col_name_lower == table_name_lower:
|
|
77
|
+
score += 4 # USER -> USER
|
|
78
|
+
break
|
|
79
|
+
|
|
80
|
+
for suffix in ['id', 'hash', 'key', 'code', 'uuid']:
|
|
81
|
+
if not col_name_lower.endswith(suffix):
|
|
82
|
+
continue
|
|
83
|
+
|
|
84
|
+
if col_name_lower == f'{table_name_lower}_{suffix}':
|
|
85
|
+
score += 5 # USER -> USER_ID
|
|
86
|
+
break
|
|
87
|
+
|
|
88
|
+
if col_name_lower == f'{table_name_lower}{suffix}':
|
|
89
|
+
score += 5 # User -> UserId
|
|
90
|
+
break
|
|
91
|
+
|
|
92
|
+
if col_name_lower.endswith(f'{table_name_lower}_{suffix}'):
|
|
93
|
+
score += 2
|
|
94
|
+
|
|
95
|
+
if col_name_lower.endswith(f'{table_name_lower}{suffix}'):
|
|
96
|
+
score += 2
|
|
97
|
+
|
|
98
|
+
# `rel-bench` hard-coding :(
|
|
99
|
+
if table_name == 'studies' and col_name == 'nct_id':
|
|
100
|
+
score += 1
|
|
101
|
+
|
|
102
|
+
ser = df[col_name].iloc[:1_000_000]
|
|
103
|
+
score += 3 * (ser.nunique() / len(ser))
|
|
104
|
+
|
|
105
|
+
scores.append((col_name, score))
|
|
106
|
+
|
|
107
|
+
scores = [x for x in scores if x[-1] >= 4]
|
|
108
|
+
scores.sort(key=lambda x: x[-1], reverse=True)
|
|
109
|
+
|
|
110
|
+
if len(scores) == 0:
|
|
111
|
+
return None
|
|
112
|
+
|
|
113
|
+
if len(scores) == 1:
|
|
114
|
+
return scores[0][0]
|
|
115
|
+
|
|
116
|
+
# In case of multiple candidates, only return one if its score is unique:
|
|
117
|
+
if scores[0][1] != scores[1][1]:
|
|
118
|
+
return scores[0][0]
|
|
119
|
+
|
|
120
|
+
max_score = max(scores, key=lambda x: x[1])
|
|
121
|
+
candidates = [col_name for col_name, score in scores if score == max_score]
|
|
122
|
+
warnings.warn(f"Found multiple potential primary keys in table "
|
|
123
|
+
f"'{table_name}': {candidates}. Please specify the primary "
|
|
124
|
+
f"key for this table manually.")
|
|
125
|
+
|
|
126
|
+
return None
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
import re
|
|
2
|
+
import warnings
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import pandas as pd
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def infer_time_column(
|
|
9
|
+
df: pd.DataFrame,
|
|
10
|
+
candidates: list[str],
|
|
11
|
+
) -> Optional[str]:
|
|
12
|
+
r"""Auto-detect potential time column.
|
|
13
|
+
|
|
14
|
+
Args:
|
|
15
|
+
df: The pandas DataFrame to analyze.
|
|
16
|
+
candidates: A list of potential candidates.
|
|
17
|
+
|
|
18
|
+
Returns:
|
|
19
|
+
The name of the detected time column, or ``None`` if not found.
|
|
20
|
+
"""
|
|
21
|
+
candidates = [ # Exclude all candidates with `*last*` in column names:
|
|
22
|
+
col_name for col_name in candidates
|
|
23
|
+
if not re.search(r'(^|_)last(_|$)', col_name, re.IGNORECASE)
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
if len(candidates) == 0:
|
|
27
|
+
return None
|
|
28
|
+
|
|
29
|
+
if len(candidates) == 1:
|
|
30
|
+
return candidates[0]
|
|
31
|
+
|
|
32
|
+
# If there exists a dedicated `create*` column, use it as time column:
|
|
33
|
+
create_candidates = [
|
|
34
|
+
candidate for candidate in candidates
|
|
35
|
+
if candidate.lower().startswith('create')
|
|
36
|
+
]
|
|
37
|
+
if len(create_candidates) == 1:
|
|
38
|
+
return create_candidates[0]
|
|
39
|
+
if len(create_candidates) > 1:
|
|
40
|
+
candidates = create_candidates
|
|
41
|
+
|
|
42
|
+
# Find the most optimal time column. Usually, it is the one pointing to
|
|
43
|
+
# the oldest timestamps:
|
|
44
|
+
with warnings.catch_warnings():
|
|
45
|
+
warnings.filterwarnings('ignore', message='Could not infer format')
|
|
46
|
+
min_timestamp_dict = {
|
|
47
|
+
key: pd.to_datetime(df[key].iloc[:10_000], 'coerce')
|
|
48
|
+
for key in candidates
|
|
49
|
+
}
|
|
50
|
+
min_timestamp_dict = {
|
|
51
|
+
key: value.min().tz_localize(None)
|
|
52
|
+
for key, value in min_timestamp_dict.items()
|
|
53
|
+
}
|
|
54
|
+
min_timestamp_dict = {
|
|
55
|
+
key: value
|
|
56
|
+
for key, value in min_timestamp_dict.items() if not pd.isna(value)
|
|
57
|
+
}
|
|
58
|
+
|
|
59
|
+
if len(min_timestamp_dict) == 0:
|
|
60
|
+
return None
|
|
61
|
+
|
|
62
|
+
return min(min_timestamp_dict, key=min_timestamp_dict.get) # type: ignore
|
|
@@ -1,3 +1,4 @@
|
|
|
1
|
+
import re
|
|
1
2
|
from typing import Dict, List, Optional, Tuple
|
|
2
3
|
|
|
3
4
|
import numpy as np
|
|
@@ -7,7 +8,47 @@ from kumoapi.typing import Stype
|
|
|
7
8
|
|
|
8
9
|
import kumoai.kumolib as kumolib
|
|
9
10
|
from kumoai.experimental.rfm.local_graph_store import LocalGraphStore
|
|
10
|
-
|
|
11
|
+
|
|
12
|
+
PUNCTUATION = re.compile(r"[\'\"\.,\(\)\!\?\;\:]")
|
|
13
|
+
MULTISPACE = re.compile(r"\s+")
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def normalize_text(
|
|
17
|
+
ser: pd.Series,
|
|
18
|
+
max_words: Optional[int] = 50,
|
|
19
|
+
) -> pd.Series:
|
|
20
|
+
r"""Normalizes text into a list of lower-case words.
|
|
21
|
+
|
|
22
|
+
Args:
|
|
23
|
+
ser: The :class:`pandas.Series` to normalize.
|
|
24
|
+
max_words: The maximum number of words to return.
|
|
25
|
+
This will auto-shrink any large text column to avoid blowing up
|
|
26
|
+
context size.
|
|
27
|
+
"""
|
|
28
|
+
if len(ser) == 0 or pd.api.types.is_list_like(ser.iloc[0]):
|
|
29
|
+
return ser
|
|
30
|
+
|
|
31
|
+
def normalize_fn(line: str) -> list[str]:
|
|
32
|
+
line = PUNCTUATION.sub(" ", line)
|
|
33
|
+
line = re.sub(r"<br\s*/?>", " ", line) # Handle <br /> or <br>
|
|
34
|
+
line = MULTISPACE.sub(" ", line)
|
|
35
|
+
words = line.split()
|
|
36
|
+
if max_words is not None:
|
|
37
|
+
words = words[:max_words]
|
|
38
|
+
return words
|
|
39
|
+
|
|
40
|
+
ser = ser.fillna('').astype(str)
|
|
41
|
+
|
|
42
|
+
if max_words is not None:
|
|
43
|
+
# We estimate the number of words as 5 characters + 1 space in an
|
|
44
|
+
# English text on average. We need this pre-filter here, as word
|
|
45
|
+
# splitting on a giant text can be very expensive:
|
|
46
|
+
ser = ser.str[:6 * max_words]
|
|
47
|
+
|
|
48
|
+
ser = ser.str.lower()
|
|
49
|
+
ser = ser.map(normalize_fn)
|
|
50
|
+
|
|
51
|
+
return ser
|
|
11
52
|
|
|
12
53
|
|
|
13
54
|
class LocalGraphSampler:
|
|
@@ -6,8 +6,7 @@ import pandas as pd
|
|
|
6
6
|
from kumoapi.rfm.context import Subgraph
|
|
7
7
|
from kumoapi.typing import Stype
|
|
8
8
|
|
|
9
|
-
from kumoai.experimental.rfm import
|
|
10
|
-
from kumoai.experimental.rfm.utils import normalize_text
|
|
9
|
+
from kumoai.experimental.rfm import Graph, LocalTable
|
|
11
10
|
from kumoai.utils import InteractiveProgressLogger, ProgressLogger
|
|
12
11
|
|
|
13
12
|
try:
|
|
@@ -20,8 +19,7 @@ except ImportError:
|
|
|
20
19
|
class LocalGraphStore:
|
|
21
20
|
def __init__(
|
|
22
21
|
self,
|
|
23
|
-
graph:
|
|
24
|
-
preprocess: bool = False,
|
|
22
|
+
graph: Graph,
|
|
25
23
|
verbose: Union[bool, ProgressLogger] = True,
|
|
26
24
|
) -> None:
|
|
27
25
|
|
|
@@ -32,7 +30,7 @@ class LocalGraphStore:
|
|
|
32
30
|
)
|
|
33
31
|
|
|
34
32
|
with verbose as logger:
|
|
35
|
-
self.df_dict, self.mask_dict = self.sanitize(graph
|
|
33
|
+
self.df_dict, self.mask_dict = self.sanitize(graph)
|
|
36
34
|
self.stype_dict = self.get_stype_dict(graph)
|
|
37
35
|
logger.log("Sanitized input data")
|
|
38
36
|
|
|
@@ -105,8 +103,7 @@ class LocalGraphStore:
|
|
|
105
103
|
|
|
106
104
|
def sanitize(
|
|
107
105
|
self,
|
|
108
|
-
graph:
|
|
109
|
-
preprocess: bool = False,
|
|
106
|
+
graph: Graph,
|
|
110
107
|
) -> Tuple[Dict[str, pd.DataFrame], Dict[str, np.ndarray]]:
|
|
111
108
|
r"""Sanitizes raw data according to table schema definition:
|
|
112
109
|
|
|
@@ -115,17 +112,12 @@ class LocalGraphStore:
|
|
|
115
112
|
* drops timezone information from timestamps
|
|
116
113
|
* drops duplicate primary keys
|
|
117
114
|
* removes rows with missing primary keys or time values
|
|
118
|
-
|
|
119
|
-
If ``preprocess`` is set to ``True``, it will additionally pre-process
|
|
120
|
-
data for faster model processing. In particular, it:
|
|
121
|
-
* tokenizes any text column that is not a foreign key
|
|
122
115
|
"""
|
|
123
|
-
df_dict: Dict[str, pd.DataFrame] = {
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
foreign_keys = {(edge.src_table, edge.fkey) for edge in graph.edges}
|
|
116
|
+
df_dict: Dict[str, pd.DataFrame] = {}
|
|
117
|
+
for table_name, table in graph.tables.items():
|
|
118
|
+
assert isinstance(table, LocalTable)
|
|
119
|
+
df = table._data
|
|
120
|
+
df_dict[table_name] = df.copy(deep=False).reset_index(drop=True)
|
|
129
121
|
|
|
130
122
|
mask_dict: Dict[str, np.ndarray] = {}
|
|
131
123
|
for table in graph.tables.values():
|
|
@@ -144,12 +136,6 @@ class LocalGraphStore:
|
|
|
144
136
|
ser = ser.dt.tz_localize(None)
|
|
145
137
|
df_dict[table.name][col.name] = ser
|
|
146
138
|
|
|
147
|
-
# Normalize text in advance (but exclude foreign keys):
|
|
148
|
-
if (preprocess and col.stype == Stype.text
|
|
149
|
-
and (table.name, col.name) not in foreign_keys):
|
|
150
|
-
ser = df_dict[table.name][col.name]
|
|
151
|
-
df_dict[table.name][col.name] = normalize_text(ser)
|
|
152
|
-
|
|
153
139
|
mask: Optional[np.ndarray] = None
|
|
154
140
|
if table._time_column is not None:
|
|
155
141
|
ser = df_dict[table.name][table._time_column]
|
|
@@ -165,7 +151,7 @@ class LocalGraphStore:
|
|
|
165
151
|
|
|
166
152
|
return df_dict, mask_dict
|
|
167
153
|
|
|
168
|
-
def get_stype_dict(self, graph:
|
|
154
|
+
def get_stype_dict(self, graph: Graph) -> Dict[str, Dict[str, Stype]]:
|
|
169
155
|
stype_dict: Dict[str, Dict[str, Stype]] = {}
|
|
170
156
|
foreign_keys = {(edge.src_table, edge.fkey) for edge in graph.edges}
|
|
171
157
|
for table in graph.tables.values():
|
|
@@ -180,7 +166,7 @@ class LocalGraphStore:
|
|
|
180
166
|
|
|
181
167
|
def get_pkey_data(
|
|
182
168
|
self,
|
|
183
|
-
graph:
|
|
169
|
+
graph: Graph,
|
|
184
170
|
) -> Tuple[
|
|
185
171
|
Dict[str, str],
|
|
186
172
|
Dict[str, pd.DataFrame],
|
|
@@ -218,7 +204,7 @@ class LocalGraphStore:
|
|
|
218
204
|
|
|
219
205
|
def get_time_data(
|
|
220
206
|
self,
|
|
221
|
-
graph:
|
|
207
|
+
graph: Graph,
|
|
222
208
|
) -> Tuple[
|
|
223
209
|
Dict[str, str],
|
|
224
210
|
Dict[str, str],
|
|
@@ -259,7 +245,7 @@ class LocalGraphStore:
|
|
|
259
245
|
|
|
260
246
|
def get_csc(
|
|
261
247
|
self,
|
|
262
|
-
graph:
|
|
248
|
+
graph: Graph,
|
|
263
249
|
) -> Tuple[
|
|
264
250
|
Dict[Tuple[str, str, str], np.ndarray],
|
|
265
251
|
Dict[Tuple[str, str, str], np.ndarray],
|
kumoai/experimental/rfm/rfm.py
CHANGED
|
@@ -32,7 +32,7 @@ from kumoapi.task import TaskType
|
|
|
32
32
|
|
|
33
33
|
from kumoai.client.rfm import RFMAPI
|
|
34
34
|
from kumoai.exceptions import HTTPException
|
|
35
|
-
from kumoai.experimental.rfm import
|
|
35
|
+
from kumoai.experimental.rfm import Graph
|
|
36
36
|
from kumoai.experimental.rfm.local_graph_sampler import LocalGraphSampler
|
|
37
37
|
from kumoai.experimental.rfm.local_graph_store import LocalGraphStore
|
|
38
38
|
from kumoai.experimental.rfm.local_pquery_driver import (
|
|
@@ -123,17 +123,17 @@ class KumoRFM:
|
|
|
123
123
|
:class:`KumoRFM` is a foundation model to generate predictions for any
|
|
124
124
|
relational dataset without training.
|
|
125
125
|
The model is pre-trained and the class provides an interface to query the
|
|
126
|
-
model from a :class:`
|
|
126
|
+
model from a :class:`Graph` object.
|
|
127
127
|
|
|
128
128
|
.. code-block:: python
|
|
129
129
|
|
|
130
|
-
from kumoai.experimental.rfm import
|
|
130
|
+
from kumoai.experimental.rfm import Graph, KumoRFM
|
|
131
131
|
|
|
132
132
|
df_users = pd.DataFrame(...)
|
|
133
133
|
df_items = pd.DataFrame(...)
|
|
134
134
|
df_orders = pd.DataFrame(...)
|
|
135
135
|
|
|
136
|
-
graph =
|
|
136
|
+
graph = Graph.from_data({
|
|
137
137
|
'users': df_users,
|
|
138
138
|
'items': df_items,
|
|
139
139
|
'orders': df_orders,
|
|
@@ -150,26 +150,16 @@ class KumoRFM:
|
|
|
150
150
|
|
|
151
151
|
Args:
|
|
152
152
|
graph: The graph.
|
|
153
|
-
preprocess: Whether to pre-process the data in advance during graph
|
|
154
|
-
materialization.
|
|
155
|
-
This is a runtime trade-off between graph materialization and model
|
|
156
|
-
processing speed.
|
|
157
|
-
It can be benefical to preprocess your data once and then run many
|
|
158
|
-
queries on top to achieve maximum model speed.
|
|
159
|
-
However, if activiated, graph materialization can take potentially
|
|
160
|
-
much longer, especially on graphs with many large text columns.
|
|
161
|
-
Best to tune this option manually.
|
|
162
153
|
verbose: Whether to print verbose output.
|
|
163
154
|
"""
|
|
164
155
|
def __init__(
|
|
165
156
|
self,
|
|
166
|
-
graph:
|
|
167
|
-
preprocess: bool = False,
|
|
157
|
+
graph: Graph,
|
|
168
158
|
verbose: Union[bool, ProgressLogger] = True,
|
|
169
159
|
) -> None:
|
|
170
160
|
graph = graph.validate()
|
|
171
161
|
self._graph_def = graph._to_api_graph_definition()
|
|
172
|
-
self._graph_store = LocalGraphStore(graph,
|
|
162
|
+
self._graph_store = LocalGraphStore(graph, verbose)
|
|
173
163
|
self._graph_sampler = LocalGraphSampler(self._graph_store)
|
|
174
164
|
|
|
175
165
|
self._client: Optional[RFMAPI] = None
|
|
@@ -2,15 +2,22 @@ import base64
|
|
|
2
2
|
import json
|
|
3
3
|
from typing import Any, Dict, List, Tuple
|
|
4
4
|
|
|
5
|
-
import boto3
|
|
6
5
|
import requests
|
|
7
|
-
from mypy_boto3_sagemaker_runtime.client import SageMakerRuntimeClient
|
|
8
|
-
from mypy_boto3_sagemaker_runtime.type_defs import InvokeEndpointOutputTypeDef
|
|
9
6
|
|
|
10
7
|
from kumoai.client import KumoClient
|
|
11
8
|
from kumoai.client.endpoints import Endpoint, HTTPMethod
|
|
12
9
|
from kumoai.exceptions import HTTPException
|
|
13
10
|
|
|
11
|
+
try:
|
|
12
|
+
# isort: off
|
|
13
|
+
from mypy_boto3_sagemaker_runtime.client import SageMakerRuntimeClient
|
|
14
|
+
from mypy_boto3_sagemaker_runtime.type_defs import (
|
|
15
|
+
InvokeEndpointOutputTypeDef, )
|
|
16
|
+
# isort: on
|
|
17
|
+
except ImportError:
|
|
18
|
+
SageMakerRuntimeClient = Any
|
|
19
|
+
InvokeEndpointOutputTypeDef = Any
|
|
20
|
+
|
|
14
21
|
|
|
15
22
|
class SageMakerResponseAdapter(requests.Response):
|
|
16
23
|
def __init__(self, sm_response: InvokeEndpointOutputTypeDef):
|
|
@@ -34,6 +41,7 @@ class SageMakerResponseAdapter(requests.Response):
|
|
|
34
41
|
|
|
35
42
|
class KumoClient_SageMakerAdapter(KumoClient):
|
|
36
43
|
def __init__(self, region: str, endpoint_name: str):
|
|
44
|
+
import boto3
|
|
37
45
|
self._client: SageMakerRuntimeClient = boto3.client(
|
|
38
46
|
service_name="sagemaker-runtime", region_name=region)
|
|
39
47
|
self._endpoint_name = endpoint_name
|
|
Binary file
|
kumoai/testing/decorators.py
CHANGED
|
@@ -25,7 +25,7 @@ def onlyFullTest(func: Callable) -> Callable:
|
|
|
25
25
|
def has_package(package: str) -> bool:
|
|
26
26
|
r"""Returns ``True`` in case ``package`` is installed."""
|
|
27
27
|
req = Requirement(package)
|
|
28
|
-
if importlib.util.find_spec(req.name) is None:
|
|
28
|
+
if importlib.util.find_spec(req.name) is None: # type: ignore
|
|
29
29
|
return False
|
|
30
30
|
|
|
31
31
|
try:
|
{kumoai-2.13.0.dev202511261731.dist-info → kumoai-2.13.0.dev202512040252.dist-info}/METADATA
RENAMED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: kumoai
|
|
3
|
-
Version: 2.13.0.
|
|
3
|
+
Version: 2.13.0.dev202512040252
|
|
4
4
|
Summary: AI on the Modern Data Stack
|
|
5
5
|
Author-email: "Kumo.AI" <hello@kumo.ai>
|
|
6
6
|
License-Expression: MIT
|
|
@@ -23,13 +23,11 @@ Requires-Dist: requests>=2.28.2
|
|
|
23
23
|
Requires-Dist: urllib3
|
|
24
24
|
Requires-Dist: plotly
|
|
25
25
|
Requires-Dist: typing_extensions>=4.5.0
|
|
26
|
-
Requires-Dist: kumo-api==0.
|
|
26
|
+
Requires-Dist: kumo-api==0.48.0
|
|
27
27
|
Requires-Dist: tqdm>=4.66.0
|
|
28
28
|
Requires-Dist: aiohttp>=3.10.0
|
|
29
29
|
Requires-Dist: pydantic>=1.10.21
|
|
30
30
|
Requires-Dist: rich>=9.0.0
|
|
31
|
-
Requires-Dist: mypy-boto3-sagemaker-runtime
|
|
32
|
-
Requires-Dist: boto3
|
|
33
31
|
Provides-Extra: doc
|
|
34
32
|
Requires-Dist: sphinx; extra == "doc"
|
|
35
33
|
Requires-Dist: sphinx-book-theme; extra == "doc"
|
|
@@ -40,13 +38,16 @@ Provides-Extra: test
|
|
|
40
38
|
Requires-Dist: pytest; extra == "test"
|
|
41
39
|
Requires-Dist: pytest-mock; extra == "test"
|
|
42
40
|
Requires-Dist: requests-mock; extra == "test"
|
|
43
|
-
Provides-Extra:
|
|
44
|
-
Requires-Dist:
|
|
45
|
-
|
|
46
|
-
Requires-Dist:
|
|
41
|
+
Provides-Extra: sqlite
|
|
42
|
+
Requires-Dist: adbc_driver_sqlite; extra == "sqlite"
|
|
43
|
+
Provides-Extra: snowflake
|
|
44
|
+
Requires-Dist: snowflake-connector-python; extra == "snowflake"
|
|
45
|
+
Requires-Dist: pyyaml; extra == "snowflake"
|
|
47
46
|
Provides-Extra: sagemaker
|
|
48
47
|
Requires-Dist: boto3<2.0,>=1.30.0; extra == "sagemaker"
|
|
49
48
|
Requires-Dist: mypy-boto3-sagemaker-runtime<2.0,>=1.34.0; extra == "sagemaker"
|
|
49
|
+
Provides-Extra: test-sagemaker
|
|
50
|
+
Requires-Dist: sagemaker<3.0; extra == "test-sagemaker"
|
|
50
51
|
Dynamic: license-file
|
|
51
52
|
Dynamic: requires-dist
|
|
52
53
|
|
|
@@ -1,13 +1,13 @@
|
|
|
1
1
|
kumoai/__init__.py,sha256=qu-qohU2cQlManX1aZIlzA3ivKl52m-cSQBPSW8urUU,10837
|
|
2
2
|
kumoai/_logging.py,sha256=qL4JbMQwKXri2f-SEJoFB8TY5ALG12S-nobGTNWxW-A,915
|
|
3
3
|
kumoai/_singleton.py,sha256=i2BHWKpccNh5SJGDyU0IXsnYzJAYr8Xb0wz4c6LRbpo,861
|
|
4
|
-
kumoai/_version.py,sha256=
|
|
4
|
+
kumoai/_version.py,sha256=16u1rVm-N2IEE7QbyS9U5nn_hjp7P_wxBIQzzAKSnDA,39
|
|
5
5
|
kumoai/databricks.py,sha256=ahwJz6DWLXMkndT0XwEDBxF-hoqhidFR8wBUQ4TLZ68,490
|
|
6
6
|
kumoai/exceptions.py,sha256=7TMs0SC8xrU009_Pgd4QXtSF9lxJq8MtRbeX9pcQUy4,859
|
|
7
7
|
kumoai/formatting.py,sha256=o3uCnLwXPhe1KI5WV9sBgRrcU7ed4rgu_pf89GL9Nc0,983
|
|
8
8
|
kumoai/futures.py,sha256=J8rtZMEYFzdn5xF_x-LAiKJz3KGL6PT02f6rq_2bOJk,3836
|
|
9
9
|
kumoai/jobs.py,sha256=dCi7BAdfm2tCnonYlGU4WJokJWbh3RzFfaOX2EYCIHU,2576
|
|
10
|
-
kumoai/kumolib.cp310-win_amd64.pyd,sha256=
|
|
10
|
+
kumoai/kumolib.cp310-win_amd64.pyd,sha256=3iE0thfrVDx0Yhh0I0li-BwZcIpQfRpaYxYMsSpYofc,194048
|
|
11
11
|
kumoai/mixin.py,sha256=IaiB8SAI0VqOoMVzzIaUlqMt53-QPUK6OB0HikG-V9E,840
|
|
12
12
|
kumoai/spcs.py,sha256=KWfENrwSLruprlD-QPh63uU0N6npiNrwkeKfBk3EUyQ,4260
|
|
13
13
|
kumoai/artifact_export/__init__.py,sha256=UXAQI5q92ChBzWAk8o3J6pElzYHudAzFZssQXd4o7i8,247
|
|
@@ -50,23 +50,35 @@ kumoai/connector/glue_connector.py,sha256=kqT2q53Da7PeeaZrvLVzFXC186E7glh5eGitKL
|
|
|
50
50
|
kumoai/connector/s3_connector.py,sha256=AUzENbQ20bYXh3XOXEOsWRKlaGGkm3YrW9JfBLm-LqY,10433
|
|
51
51
|
kumoai/connector/snowflake_connector.py,sha256=tQzIWxC4oDGqxFt0212w5eoIPT4QBP2nuF9SdKRNwNI,9274
|
|
52
52
|
kumoai/connector/source_table.py,sha256=fnqwIKY6qYo4G0EsRzchb6FgZ-dQyU6aRaD9UAxsml0,18010
|
|
53
|
-
kumoai/connector/utils.py,sha256=
|
|
53
|
+
kumoai/connector/utils.py,sha256=5K9BMdWiIP3hhdkUc6Xt1e0xv5YyziXtZ4PnBqq0Ehw,66490
|
|
54
54
|
kumoai/encoder/__init__.py,sha256=8FeP6mUyCeXxr1b8kUIi5dxe5vEXQRft9tPoaV1CBqg,186
|
|
55
55
|
kumoai/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
56
|
-
kumoai/experimental/rfm/__init__.py,sha256=
|
|
56
|
+
kumoai/experimental/rfm/__init__.py,sha256=EFZz6IvvskmeO85Vig6p1m_6jdimS_BkeREOndHuRsc,6247
|
|
57
57
|
kumoai/experimental/rfm/authenticate.py,sha256=G89_4TMeUpr5fG_0VTzMF5sdNhaciitA1oc2loTlTmo,19321
|
|
58
|
-
kumoai/experimental/rfm/
|
|
59
|
-
kumoai/experimental/rfm/local_graph_sampler.py,sha256=
|
|
60
|
-
kumoai/experimental/rfm/local_graph_store.py,sha256=
|
|
58
|
+
kumoai/experimental/rfm/graph.py,sha256=kSWve-Fn_9qERFjEpCDO5zDnngtd9T4MOhR_o46PI7s,39602
|
|
59
|
+
kumoai/experimental/rfm/local_graph_sampler.py,sha256=dQ3JnuozTNeZyUFRu2h8OTMNmV1RAoaCA0gvkpgOstg,8110
|
|
60
|
+
kumoai/experimental/rfm/local_graph_store.py,sha256=6jY1ciVIlnBBhZCxWwBTl7SKX1fxRIDLszwrftD0Cdk,13485
|
|
61
61
|
kumoai/experimental/rfm/local_pquery_driver.py,sha256=Yd_yHIrvuDj16IC1pvsqiQvZS41vvOOCRMiuDGtN6Fk,26851
|
|
62
|
-
kumoai/experimental/rfm/
|
|
63
|
-
kumoai/experimental/rfm/
|
|
64
|
-
kumoai/experimental/rfm/
|
|
65
|
-
kumoai/experimental/rfm/
|
|
66
|
-
kumoai/experimental/rfm/
|
|
62
|
+
kumoai/experimental/rfm/rfm.py,sha256=vOnL8ecHTo1TX2B8_T8xaWGou8qYYz8DyVENu1H93mM,48834
|
|
63
|
+
kumoai/experimental/rfm/sagemaker.py,sha256=sEJSyfEFBA3-7wKinBEzSooKHEn0BgPjrgRnPhYo79g,5120
|
|
64
|
+
kumoai/experimental/rfm/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
65
|
+
kumoai/experimental/rfm/backend/local/__init__.py,sha256=usMh0fuDxKK-aOVT1sU30BQWFS0eSkfUrhUVILisQQI,934
|
|
66
|
+
kumoai/experimental/rfm/backend/local/table.py,sha256=1PqNOROzlnK3SaZHNcU2hyzeifs0N4wssQAS3-Z0Myc,3674
|
|
67
|
+
kumoai/experimental/rfm/backend/snow/__init__.py,sha256=viMeR9VWpB1kjRdSWCTNFMdM7a8Mj_Dtck1twJW8dV8,962
|
|
68
|
+
kumoai/experimental/rfm/backend/snow/table.py,sha256=HVrPtCVvfsisFmq9jMovowsE5Wl5oti3O-kru7ruXlc,4312
|
|
69
|
+
kumoai/experimental/rfm/backend/sqlite/__init__.py,sha256=xw5NNLrWSvUvRkD49X_9hZYjas5EuP1XDANPy0EEjOg,874
|
|
70
|
+
kumoai/experimental/rfm/backend/sqlite/table.py,sha256=mBiZC21gQwfR4demFrP37GmawMHfIm-G82mLQeBqIZo,3901
|
|
71
|
+
kumoai/experimental/rfm/base/__init__.py,sha256=oXPkeBemtuDxRUK61-0sOT84GZB_oQ6HvaZNU1KFNaw,199
|
|
72
|
+
kumoai/experimental/rfm/base/column.py,sha256=OE-PRQ8HO4uTq0e3_3eHJFfhp5nzw79zd-43g3iMh4g,2385
|
|
73
|
+
kumoai/experimental/rfm/base/source.py,sha256=H5yN9xAwK3i_69EdqOV_x58muPGKQiI8ev5BhHQDZEo,290
|
|
74
|
+
kumoai/experimental/rfm/base/table.py,sha256=glyAg4LCQdddM3lIRClJSA7qMyfoHUVAGBf1rEs6B8Y,20113
|
|
75
|
+
kumoai/experimental/rfm/infer/__init__.py,sha256=qKg8or-SpgTApD6ePw1PJ4aUZPrOLTHLRCmBIJ92hrk,486
|
|
67
76
|
kumoai/experimental/rfm/infer/categorical.py,sha256=bqmfrE5ZCBTcb35lA4SyAkCu3MgttAn29VBJYMBNhVg,893
|
|
77
|
+
kumoai/experimental/rfm/infer/dtype.py,sha256=Hf_drluYNuN59lTSe-8GuXalg20Pv93kCktB6Hb9f74,2686
|
|
68
78
|
kumoai/experimental/rfm/infer/id.py,sha256=xaJBETLZa8ttzZCsDwFSwfyCi3VYsLc_kDWT_t_6Ih4,954
|
|
69
79
|
kumoai/experimental/rfm/infer/multicategorical.py,sha256=D-1KwYRkOSkBrOJr4Xa3eTCoAF9O9hPGa7Vg67V5_HU,1150
|
|
80
|
+
kumoai/experimental/rfm/infer/pkey.py,sha256=Hvztcircd4iGdsnFU9Xi1kq_A5ONMnkAdnrpQT5svSs,4519
|
|
81
|
+
kumoai/experimental/rfm/infer/time_col.py,sha256=G98Cgz1m9G9VA-ApnCmGYnJxEFwp1jfaPf3nCMOz_N0,1882
|
|
70
82
|
kumoai/experimental/rfm/infer/timestamp.py,sha256=L2VxjtYTSyUBYAo4M-L08xSQlPpqnHMAVF5_vxjh3Y0,1135
|
|
71
83
|
kumoai/experimental/rfm/pquery/__init__.py,sha256=RkTn0I74uXOUuOiBpa6S-_QEYctMutkUnBEfF9ztQzI,159
|
|
72
84
|
kumoai/experimental/rfm/pquery/executor.py,sha256=S8wwXbAkH-YSnmEVYB8d6wyJF4JJ003mH_0zFTvOp_I,2843
|
|
@@ -80,7 +92,7 @@ kumoai/pquery/prediction_table.py,sha256=hWG4L_ze4PLgUoxCXNKk8_nkYxVXELQs8_X8KGO
|
|
|
80
92
|
kumoai/pquery/predictive_query.py,sha256=GWhQpQxf6apyyu-bvE3z63mX6NLd8lKbyu_jzj7rNms,25608
|
|
81
93
|
kumoai/pquery/training_table.py,sha256=L1QjaVlY4SAPD8OUmTaH6YjZzBbPOnS9mnAT69znWv0,16233
|
|
82
94
|
kumoai/testing/__init__.py,sha256=XBQ_Sa3WnOYlpXZ3gUn8w6nVfZt-nfPhytfIBeiPt4w,178
|
|
83
|
-
kumoai/testing/decorators.py,sha256=
|
|
95
|
+
kumoai/testing/decorators.py,sha256=p79ZCQqPY_MHWy0_l7-xQ6wUIqFTn4AbrGWTHLvpbQY,1664
|
|
84
96
|
kumoai/trainer/__init__.py,sha256=uCFXy9bw_byn_wYd3M-BTZCHTVvv4XXr8qRlh-QOvag,981
|
|
85
97
|
kumoai/trainer/baseline_trainer.py,sha256=oXweh8j1sar6KhQfr3A7gmQxcDq7SG0Bx3jIenbtyC4,4117
|
|
86
98
|
kumoai/trainer/config.py,sha256=7_Jv1w1mqaokCQwQdJkqCSgVpmh8GqE3fL1Ky_vvttI,100
|
|
@@ -92,8 +104,8 @@ kumoai/utils/__init__.py,sha256=wAKgmwtMIGuiauW9D_GGKH95K-24Kgwmld27mm4nsro,278
|
|
|
92
104
|
kumoai/utils/datasets.py,sha256=UyAII-oAn7x3ombuvpbSQ41aVF9SYKBjQthTD-vcT2A,3011
|
|
93
105
|
kumoai/utils/forecasting.py,sha256=ZgKeUCbWLOot0giAkoigwU5du8LkrwAicFOi5hVn6wg,7624
|
|
94
106
|
kumoai/utils/progress_logger.py,sha256=MZsWgHd4UZQKCXiJZgQeW-Emi_BmzlCKPLPXOL_HqBo,5239
|
|
95
|
-
kumoai-2.13.0.
|
|
96
|
-
kumoai-2.13.0.
|
|
97
|
-
kumoai-2.13.0.
|
|
98
|
-
kumoai-2.13.0.
|
|
99
|
-
kumoai-2.13.0.
|
|
107
|
+
kumoai-2.13.0.dev202512040252.dist-info/licenses/LICENSE,sha256=ZUilBDp--4vbhsEr6f_Upw9rnIx09zQ3K9fXQ0rfd6w,1111
|
|
108
|
+
kumoai-2.13.0.dev202512040252.dist-info/METADATA,sha256=T-O--qEm_2QPzB-dDkwR6Ei7r79H7v6qQBbR1e1J8gg,2580
|
|
109
|
+
kumoai-2.13.0.dev202512040252.dist-info/WHEEL,sha256=KUuBC6lxAbHCKilKua8R9W_TM71_-9Sg5uEP3uDWcoU,101
|
|
110
|
+
kumoai-2.13.0.dev202512040252.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
|
|
111
|
+
kumoai-2.13.0.dev202512040252.dist-info/RECORD,,
|