dpst 0.1.0__py3-none-any.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.
- dpst/__init__.py +5 -0
- dpst/cli.py +27 -0
- dpst/core.py +277 -0
- dpst/setup_db.py +156 -0
- dpst-0.1.0.dist-info/METADATA +132 -0
- dpst-0.1.0.dist-info/RECORD +10 -0
- dpst-0.1.0.dist-info/WHEEL +5 -0
- dpst-0.1.0.dist-info/entry_points.txt +2 -0
- dpst-0.1.0.dist-info/licenses/LICENSE +21 -0
- dpst-0.1.0.dist-info/top_level.txt +1 -0
dpst/__init__.py
ADDED
dpst/cli.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
# src/dpst/cli.py
|
|
2
|
+
import sys
|
|
3
|
+
import argparse
|
|
4
|
+
from dpst.setup_db import initialize_database, run_clustering_and_indexing
|
|
5
|
+
|
|
6
|
+
def main():
|
|
7
|
+
parser = argparse.ArgumentParser(description="DPST Package Management Utility")
|
|
8
|
+
subparsers = parser.add_subparsers(dest="command")
|
|
9
|
+
|
|
10
|
+
subparsers.add_parser("setup", help="Initialize Weaviate DB, download FineWeb, and generate cluster JSON metadata files.")
|
|
11
|
+
|
|
12
|
+
args = parser.parse_args()
|
|
13
|
+
|
|
14
|
+
if args.command == "setup":
|
|
15
|
+
print("--- Starting DPST Setup Process ---")
|
|
16
|
+
try:
|
|
17
|
+
initialize_database()
|
|
18
|
+
run_clustering_and_indexing()
|
|
19
|
+
print("--- Setup Successful! You can now use DPST ---")
|
|
20
|
+
except Exception as e:
|
|
21
|
+
print(f"Setup failed with error: {e}", file=sys.stderr)
|
|
22
|
+
sys.exit(1)
|
|
23
|
+
else:
|
|
24
|
+
parser.print_help()
|
|
25
|
+
|
|
26
|
+
if __name__ == "__main__":
|
|
27
|
+
main()
|
dpst/core.py
ADDED
|
@@ -0,0 +1,277 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import json
|
|
3
|
+
from collections import defaultdict
|
|
4
|
+
from functools import partial
|
|
5
|
+
import importlib.resources as pkg_resources
|
|
6
|
+
|
|
7
|
+
import numpy as np
|
|
8
|
+
import torch
|
|
9
|
+
from torch.nn import CrossEntropyLoss
|
|
10
|
+
from tqdm.auto import tqdm
|
|
11
|
+
from datasketch import MinHash, MinHashLSH
|
|
12
|
+
from nltk import ngrams
|
|
13
|
+
|
|
14
|
+
import weaviate
|
|
15
|
+
from weaviate.classes.query import MetadataQuery, Filter
|
|
16
|
+
|
|
17
|
+
from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, pipeline
|
|
18
|
+
from sentence_transformers import util
|
|
19
|
+
from openie import StanfordOpenIE
|
|
20
|
+
|
|
21
|
+
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
|
|
22
|
+
|
|
23
|
+
class DPST:
|
|
24
|
+
def __init__(self, mode: str, hf_token: str = None, model_checkpoint: str = "meta-llama/Llama-3.2-1B-Instruct"):
|
|
25
|
+
"""
|
|
26
|
+
DPST Runtime Engine. Assumes that `dpst setup` has already been executed
|
|
27
|
+
to spin up the Weaviate collections and generate the local data clusters.
|
|
28
|
+
"""
|
|
29
|
+
print("Initializing DPST Runtime...", flush=True)
|
|
30
|
+
|
|
31
|
+
mode_map = {
|
|
32
|
+
"50k": "fiftyk",
|
|
33
|
+
"100k": "hundredk",
|
|
34
|
+
"200k": "twohundredk"
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
if mode not in mode_map:
|
|
38
|
+
raise ValueError("Error: [MODE] must be one of [50k, 100k, 200k].")
|
|
39
|
+
|
|
40
|
+
self.mode = mode_map[mode]
|
|
41
|
+
|
|
42
|
+
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
43
|
+
|
|
44
|
+
try:
|
|
45
|
+
data_package = pkg_resources.files("dpst.data")
|
|
46
|
+
|
|
47
|
+
with open(data_package / f"{mode}.json", 'r') as f:
|
|
48
|
+
self.centroids = torch.tensor(json.load(f)).to(self.device)
|
|
49
|
+
|
|
50
|
+
with open(data_package / f"{mode}_counts.json", 'r') as f:
|
|
51
|
+
self.cluster_counts = json.load(f)
|
|
52
|
+
except FileNotFoundError:
|
|
53
|
+
raise RuntimeError(
|
|
54
|
+
f"Cluster data files for mode '{mode}' were not found. "
|
|
55
|
+
"Please run `dpst setup` in your terminal to initialize database assets before using the runtime."
|
|
56
|
+
)
|
|
57
|
+
|
|
58
|
+
self.client = weaviate.connect_to_local()
|
|
59
|
+
try:
|
|
60
|
+
self.collection = self.client.collections.get("Triples")
|
|
61
|
+
except Exception:
|
|
62
|
+
self.client.close()
|
|
63
|
+
raise RuntimeError(
|
|
64
|
+
"The Weaviate collection 'Triples' does not exist. "
|
|
65
|
+
"Ensure your local Weaviate instance is running and you have executed `dpst setup`."
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
self.properties = {
|
|
69
|
+
"openie.affinity_probability_cap": 2 / 3,
|
|
70
|
+
"openie.triple.strict": False,
|
|
71
|
+
}
|
|
72
|
+
self.IEclient = StanfordOpenIE(properties=self.properties)
|
|
73
|
+
|
|
74
|
+
self.model_checkpoint = model_checkpoint
|
|
75
|
+
self.model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code=True).to(self.device)
|
|
76
|
+
|
|
77
|
+
self.gen_model = AutoModelForCausalLM.from_pretrained(self.model_checkpoint, token=hf_token, device_map=self.device)
|
|
78
|
+
self.tokenizer = AutoTokenizer.from_pretrained(self.model_checkpoint, token=hf_token)
|
|
79
|
+
|
|
80
|
+
self.pipe = pipeline(
|
|
81
|
+
"text-generation",
|
|
82
|
+
model=self.gen_model,
|
|
83
|
+
tokenizer=self.tokenizer,
|
|
84
|
+
torch_dtype=torch.bfloat16,
|
|
85
|
+
device_map="auto",
|
|
86
|
+
)
|
|
87
|
+
|
|
88
|
+
self.ppl_model = AutoModelForCausalLM.from_pretrained("gpt2").to(self.device)
|
|
89
|
+
self.ppl_tokenizer = AutoTokenizer.from_pretrained("gpt2")
|
|
90
|
+
|
|
91
|
+
print("Initialization Finished.", flush=True)
|
|
92
|
+
|
|
93
|
+
def cleanup(self):
|
|
94
|
+
if hasattr(self, 'client') and self.client:
|
|
95
|
+
self.client.close()
|
|
96
|
+
|
|
97
|
+
def exponential(self, candidates, epsilon, sensitivity=1):
|
|
98
|
+
probabilities = [np.exp(epsilon * x[1] / (2 * sensitivity)) for x in candidates]
|
|
99
|
+
probabilities = probabilities / np.linalg.norm(probabilities, ord=1)
|
|
100
|
+
return np.random.choice([x[0] for x in candidates], 1, p=probabilities)[0]
|
|
101
|
+
|
|
102
|
+
def query_db(self, vector, cluster):
|
|
103
|
+
response = self.collection.query.near_vector(
|
|
104
|
+
near_vector=vector,
|
|
105
|
+
limit=self.cluster_counts[cluster],
|
|
106
|
+
filters=Filter.by_property(self.mode).equal(cluster),
|
|
107
|
+
return_metadata=MetadataQuery(distance=True)
|
|
108
|
+
)
|
|
109
|
+
candidates = [(x.properties["text"], max(1 - x.metadata.distance, 0)) for x in response.objects]
|
|
110
|
+
return candidates
|
|
111
|
+
|
|
112
|
+
def get_prompt(self, triples, messages=True):
|
|
113
|
+
PROMPT = [
|
|
114
|
+
{"role": "system", "content": "Generate a concise text for the given set of triples. Ensure that the generated output only includes the provided information from the triples, but feel free to fill in the gaps where sensible. If necessary, ignore triples that do not fit into the larger context. It is very important that the output is grammatically correct, natural, and logical. Provide a text that captures the semantic meaning of the triples, without being too verbose or lengthy. Do not provide any further explanation, only provide the output text."},
|
|
115
|
+
{"role": "user", "content": "Input triples: [{’object’: ’Mike_Mularkey’,’property’: ’coach’,’subject’: ’Tennessee_Titans’}]"},
|
|
116
|
+
{"role": "assistant", "content": "Output text: Mike Mularkey is the coach of the Tennessee Titans."},
|
|
117
|
+
{"role": "user", "content": "Input triples: [{’object’: ’Albert_E._Austin’, ’property’: ’successor’, ’subject’: ’Alfred_N._Phillips’}, {’object’: ’Connecticut’, ’property’: ’birthPlace’, ’subject’: ’Alfred_N._Phillips’}, {’object’: ’United_States_House_of_Representatives’, ’property’: ’office’, ’subject’: ’Alfred_N._Phillips’}]"},
|
|
118
|
+
{"role": "assistant", "content": "Output text: Albert E. Austin succeeded Alfred N. Phillips who was born in Connecticut and worked at the United States House of Representatives."},
|
|
119
|
+
{"role": "user", "content": "Input triples: [{’object’: ’College_of_William_&_Mary’, ’property’: ’owner’, ’subject’: ’Alan_B._Miller_Hall’}, {’object’: ’2009-06-01’, ’property’: ’completionDate’, ’subject’: ’Alan_B._Miller_Hall’}, {’object’: ’101 Ukrop Way’, ’property’: ’address’, ’subject’: ’Alan_B._Miller_Hall’}, {’object’: ’Williamsburg,_Virginia’, ’property’: ’location’, ’subject’: ’Alan_B._Miller_Hall’}, {’object’: ’Robert_A._M._Stern’, ’property’: ’architect’, ’subject’: ’Alan_B._Miller_Hall’}]"},
|
|
120
|
+
{"role": "assistant", "content": "Output text: The Alan B Miller Hall’s location is 101 Ukrop Way, Williamsburg, Virginia. It was designed by Robert A.M. Stern and was completed on 1 June 2009. Its owner is the College of William and Mary."},
|
|
121
|
+
{"role": "user", "content": "Input Triples: {}".format(str(triples))}
|
|
122
|
+
]
|
|
123
|
+
return PROMPT
|
|
124
|
+
|
|
125
|
+
def compute_ppl(self, predictions, batch_size: int = 16, add_start_token: bool = True, max_length=32):
|
|
126
|
+
if self.ppl_tokenizer.pad_token is None and batch_size > 1:
|
|
127
|
+
existing_special_tokens = list(self.ppl_tokenizer.special_tokens_map_extended.values())
|
|
128
|
+
assert len(existing_special_tokens) > 0, "Model must have at least one special token to use for padding."
|
|
129
|
+
self.ppl_tokenizer.add_special_tokens({"pad_token": existing_special_tokens[0]})
|
|
130
|
+
|
|
131
|
+
max_tokenized_len = max_length - 1 if (add_start_token and max_length) else max_length
|
|
132
|
+
|
|
133
|
+
encodings = self.ppl_tokenizer(
|
|
134
|
+
predictions,
|
|
135
|
+
add_special_tokens=False,
|
|
136
|
+
padding=True,
|
|
137
|
+
truncation=True if max_tokenized_len else False,
|
|
138
|
+
max_length=max_tokenized_len,
|
|
139
|
+
return_tensors="pt",
|
|
140
|
+
return_attention_mask=True,
|
|
141
|
+
).to(self.device)
|
|
142
|
+
|
|
143
|
+
encoded_texts = encodings["input_ids"]
|
|
144
|
+
attn_masks = encodings["attention_mask"]
|
|
145
|
+
|
|
146
|
+
if add_start_token:
|
|
147
|
+
assert torch.all(torch.ge(attn_masks.sum(1), 1)), "Each input text must be at least one token long."
|
|
148
|
+
else:
|
|
149
|
+
assert torch.all(torch.ge(attn_masks.sum(1), 2)), "Each input text must be at least two tokens long."
|
|
150
|
+
|
|
151
|
+
ppls = []
|
|
152
|
+
loss_fct = CrossEntropyLoss(reduction="none")
|
|
153
|
+
|
|
154
|
+
for start_index in range(0, len(encoded_texts), batch_size):
|
|
155
|
+
end_index = min(start_index + batch_size, len(encoded_texts))
|
|
156
|
+
encoded_batch = encoded_texts[start_index:end_index]
|
|
157
|
+
attn_mask = attn_masks[start_index:end_index]
|
|
158
|
+
|
|
159
|
+
if add_start_token:
|
|
160
|
+
bos_tokens_tensor = torch.tensor([[self.ppl_tokenizer.bos_token_id]] * encoded_batch.size(dim=0)).to(self.device)
|
|
161
|
+
encoded_batch = torch.cat([bos_tokens_tensor, encoded_batch], dim=1)
|
|
162
|
+
attn_mask = torch.cat([torch.ones(bos_tokens_tensor.size(), dtype=torch.int64).to(self.device), attn_mask], dim=1)
|
|
163
|
+
|
|
164
|
+
labels = encoded_batch
|
|
165
|
+
|
|
166
|
+
with torch.no_grad():
|
|
167
|
+
out_logits = self.ppl_model(encoded_batch, attention_mask=attn_mask).logits
|
|
168
|
+
|
|
169
|
+
shift_logits = out_logits[..., :-1, :].contiguous()
|
|
170
|
+
shift_labels = labels[..., 1:].contiguous()
|
|
171
|
+
shift_attention_mask_batch = attn_mask[..., 1:].contiguous()
|
|
172
|
+
|
|
173
|
+
perplexity_batch = torch.exp(
|
|
174
|
+
(loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).sum(1)
|
|
175
|
+
/ shift_attention_mask_batch.sum(1)
|
|
176
|
+
)
|
|
177
|
+
ppls += perplexity_batch.tolist()
|
|
178
|
+
|
|
179
|
+
return {"perplexities": ppls, "mean_perplexity": np.mean(ppls)}
|
|
180
|
+
|
|
181
|
+
def get_triples_ie(self, text):
|
|
182
|
+
res = [x for x in self.IEclient.annotate(text)]
|
|
183
|
+
temp = [tuple(x.values()) for x in res]
|
|
184
|
+
|
|
185
|
+
current = defaultdict(list)
|
|
186
|
+
for t in temp:
|
|
187
|
+
current[(t[0], t[1])].append(t)
|
|
188
|
+
|
|
189
|
+
final = []
|
|
190
|
+
for t in temp:
|
|
191
|
+
s = "{} | {} | {}".format(t[0], t[1], t[2])
|
|
192
|
+
if s not in final:
|
|
193
|
+
final.append(s.replace("_", " "))
|
|
194
|
+
|
|
195
|
+
lsh = MinHashLSH(threshold=0.4, num_perm=128)
|
|
196
|
+
minhashes = {}
|
|
197
|
+
for i, f in enumerate(final):
|
|
198
|
+
minhash = MinHash(num_perm=128)
|
|
199
|
+
for d in ngrams(f, 3):
|
|
200
|
+
minhash.update("".join(d).encode('utf-8'))
|
|
201
|
+
lsh.insert(i, minhash)
|
|
202
|
+
minhashes[i] = minhash
|
|
203
|
+
|
|
204
|
+
matches = {}
|
|
205
|
+
for x, y in zip(final, minhashes):
|
|
206
|
+
matches[x] = [final[z] for z in lsh.query(minhashes[y]) if z != y]
|
|
207
|
+
|
|
208
|
+
clusters = []
|
|
209
|
+
covered = []
|
|
210
|
+
for m in sorted(matches, key=lambda x: len(matches[x]), reverse=True):
|
|
211
|
+
if m not in covered and len(matches[m]) > 0:
|
|
212
|
+
clusters.append(matches[m])
|
|
213
|
+
covered.extend(matches[m])
|
|
214
|
+
|
|
215
|
+
clean = [x.replace(" | ", " ") for x in covered]
|
|
216
|
+
if len(clean) == 0:
|
|
217
|
+
return []
|
|
218
|
+
ppls = dict(zip(covered, self.compute_ppl(predictions=clean, batch_size=64)["perplexities"]))
|
|
219
|
+
|
|
220
|
+
best = []
|
|
221
|
+
for c in clusters:
|
|
222
|
+
scores = [ppls[x] for x in c]
|
|
223
|
+
imin = np.argmin(scores)
|
|
224
|
+
best.append(c[imin])
|
|
225
|
+
|
|
226
|
+
ordered = []
|
|
227
|
+
for f in final:
|
|
228
|
+
if f in best:
|
|
229
|
+
ordered.append(f)
|
|
230
|
+
return ordered
|
|
231
|
+
|
|
232
|
+
def privatize(self, texts, epsilon=10, DP=True):
|
|
233
|
+
results = []
|
|
234
|
+
for i, t in tqdm(enumerate(texts), total=len(texts)):
|
|
235
|
+
triples = self.get_triples_ie(t)
|
|
236
|
+
if len(triples) == 0:
|
|
237
|
+
results.append(t)
|
|
238
|
+
continue
|
|
239
|
+
|
|
240
|
+
if DP:
|
|
241
|
+
eps = epsilon / len(triples)
|
|
242
|
+
query_vectors = self.model.encode(triples, task="text-matching", truncate_dim=32, max_length=64)
|
|
243
|
+
|
|
244
|
+
res = util.semantic_search(query_embeddings=torch.tensor(query_vectors).to(self.device), corpus_embeddings=self.centroids, top_k=1)
|
|
245
|
+
clusters = [r[0]["corpus_id"] for r in res]
|
|
246
|
+
|
|
247
|
+
candidates = []
|
|
248
|
+
for q, c in zip(query_vectors, clusters):
|
|
249
|
+
near = self.query_db(q, c)
|
|
250
|
+
if len(near) > 0:
|
|
251
|
+
candidates.append(near)
|
|
252
|
+
private_triples = [self.exponential(c, eps) for c in candidates]
|
|
253
|
+
|
|
254
|
+
final = []
|
|
255
|
+
for p in private_triples:
|
|
256
|
+
m = p.split(" | ")
|
|
257
|
+
final.append({"object": m[0], "property": m[1], "subject": m[2]})
|
|
258
|
+
prompt = self.get_prompt(final)
|
|
259
|
+
else:
|
|
260
|
+
final = []
|
|
261
|
+
for x in triples:
|
|
262
|
+
m = x.split(" | ")
|
|
263
|
+
final.append({"object": m[0], "property": m[1], "subject": m[2]})
|
|
264
|
+
prompt = self.get_prompt(final)
|
|
265
|
+
|
|
266
|
+
outputs = self.pipe(
|
|
267
|
+
prompt,
|
|
268
|
+
pad_token_id=self.tokenizer.eos_token_id,
|
|
269
|
+
max_new_tokens=int(len(self.tokenizer.encode(texts[i], return_tensors="pt")[0]))
|
|
270
|
+
)
|
|
271
|
+
generated = outputs[0]["generated_text"][-1]["content"]
|
|
272
|
+
|
|
273
|
+
generated = generated.split("Output text: ")[-1].strip().replace("\n", "")
|
|
274
|
+
generated = generated.split("USER:")[0].strip().replace("\n", "")
|
|
275
|
+
generated = generated.split("\t")[0].split("ASSISTANT")[0].split("USER")[0].split("###")[0].split("Note:")[0].split("Explanation:")[0].split("```")[0].split("EXPECTED_OUTPUT")[0]
|
|
276
|
+
results.append(generated.strip())
|
|
277
|
+
return results
|
dpst/setup_db.py
ADDED
|
@@ -0,0 +1,156 @@
|
|
|
1
|
+
# src/dpst/setup_db.py
|
|
2
|
+
import json
|
|
3
|
+
import numpy as np
|
|
4
|
+
from tqdm.auto import tqdm
|
|
5
|
+
import weaviate
|
|
6
|
+
import weaviate.classes.config as wc
|
|
7
|
+
from weaviate.util import generate_uuid5
|
|
8
|
+
from importlib import resources
|
|
9
|
+
from collections import Counter
|
|
10
|
+
|
|
11
|
+
def initialize_database(max_rows=100000):
|
|
12
|
+
"""Create the DB collection, parse FineWeb text, embed, and store."""
|
|
13
|
+
from datasets import load_dataset
|
|
14
|
+
|
|
15
|
+
print("Connecting to local Weaviate instance...")
|
|
16
|
+
client = weaviate.connect_to_local()
|
|
17
|
+
|
|
18
|
+
if not client.collections.exists("Triples"):
|
|
19
|
+
print("Creating 'Triples' collection schema...")
|
|
20
|
+
client.collections.create(
|
|
21
|
+
name="Triples",
|
|
22
|
+
properties=[
|
|
23
|
+
wc.Property(name="text", data_type=wc.DataType.TEXT),
|
|
24
|
+
wc.Property(name="fiftyk", data_type=wc.DataType.INT),
|
|
25
|
+
wc.Property(name="hundredk", data_type=wc.DataType.INT),
|
|
26
|
+
wc.Property(name="twohundredk", data_type=wc.DataType.INT),
|
|
27
|
+
],
|
|
28
|
+
vectorizer_config=wc.Configure.Vectorizer.none(),
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
triples = client.collections.get("Triples")
|
|
32
|
+
|
|
33
|
+
from dpst.core import embedding_model, get_triples_ie
|
|
34
|
+
|
|
35
|
+
print("Streaming sample-10BT from HuggingFaceFW/fineweb...")
|
|
36
|
+
dataset = load_dataset("HuggingFaceFW/fineweb", "sample-10BT", split="train", streaming=True)
|
|
37
|
+
|
|
38
|
+
all_data = []
|
|
39
|
+
total = 0
|
|
40
|
+
|
|
41
|
+
for row in dataset:
|
|
42
|
+
if idx == max_rows:
|
|
43
|
+
break
|
|
44
|
+
|
|
45
|
+
if idx % 100000 == 0 and idx != 0:
|
|
46
|
+
print("{} rows processed.".format(idx))
|
|
47
|
+
|
|
48
|
+
if len(row["text"]) > 10000 or row["language"] != "en":
|
|
49
|
+
continue
|
|
50
|
+
|
|
51
|
+
all_data.append(row["text"])
|
|
52
|
+
|
|
53
|
+
if len(all_data) >= 1000:
|
|
54
|
+
print("Extracting triples...")
|
|
55
|
+
res = [get_triples_ie(x) for x in tqdm(all_data)]
|
|
56
|
+
tri = []
|
|
57
|
+
for r in res:
|
|
58
|
+
tri.extend(r)
|
|
59
|
+
|
|
60
|
+
total += len(tri)
|
|
61
|
+
print("Inserting {} triples...".format(len(tri)))
|
|
62
|
+
embeddings = embedding_model.encode(tri, task="text-matching", truncate_dim=32, max_length=64)
|
|
63
|
+
|
|
64
|
+
with triples.batch.dynamic() as batch:
|
|
65
|
+
for i, t in enumerate(tri):
|
|
66
|
+
obj = {"text":t}
|
|
67
|
+
vector = embeddings[i]
|
|
68
|
+
|
|
69
|
+
batch.add_object(
|
|
70
|
+
properties=obj,
|
|
71
|
+
uuid=generate_uuid5(obj),
|
|
72
|
+
vector=vector
|
|
73
|
+
)
|
|
74
|
+
print("Finished. Total: {}".format(total))
|
|
75
|
+
del tri[:]
|
|
76
|
+
del tri
|
|
77
|
+
del all_data[:]
|
|
78
|
+
all_data = []
|
|
79
|
+
|
|
80
|
+
idx += 1
|
|
81
|
+
|
|
82
|
+
print("Extracting triples...")
|
|
83
|
+
res = [get_triples_ie(x) for x in tqdm(all_data)]
|
|
84
|
+
tri = []
|
|
85
|
+
for r in res:
|
|
86
|
+
tri.extend(r)
|
|
87
|
+
|
|
88
|
+
total += len(tri)
|
|
89
|
+
print("Inserting {} triples...".format(len(tri)))
|
|
90
|
+
embeddings = embedding_model.encode(tri, task="text-matching", truncate_dim=32, max_length=64)
|
|
91
|
+
|
|
92
|
+
with triples.batch.dynamic() as batch:
|
|
93
|
+
for i, t in enumerate(tri):
|
|
94
|
+
obj = {"text":t}
|
|
95
|
+
vector = embeddings[i]
|
|
96
|
+
|
|
97
|
+
batch.add_object(
|
|
98
|
+
properties=obj,
|
|
99
|
+
uuid=generate_uuid5(obj),
|
|
100
|
+
vector=vector
|
|
101
|
+
)
|
|
102
|
+
print("Finished. Total: {}".format(total))
|
|
103
|
+
del tri[:]
|
|
104
|
+
del tri
|
|
105
|
+
del all_data[:]
|
|
106
|
+
all_data = []
|
|
107
|
+
|
|
108
|
+
client.close()
|
|
109
|
+
print("Database seeding completed.")
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def run_clustering_and_indexing():
|
|
113
|
+
"""Pull vectors, run K-Means, export cluster JSONs, and update Weaviate."""
|
|
114
|
+
from sklearn.cluster import MiniBatchKMeans
|
|
115
|
+
|
|
116
|
+
client = weaviate.connect_to_local()
|
|
117
|
+
triples = client.collections.get("Triples")
|
|
118
|
+
|
|
119
|
+
print("Fetching vectors from Weaviate...")
|
|
120
|
+
all_ids, all_vectors, all_texts = [], [], []
|
|
121
|
+
for t in tqdm(triples.iterator(include_vector=True)):
|
|
122
|
+
all_ids.append(t.uuid)
|
|
123
|
+
all_vectors.append(t.vector["default"])
|
|
124
|
+
|
|
125
|
+
x = np.array(all_vectors)
|
|
126
|
+
|
|
127
|
+
print("Running MiniBatchKMeans (50k, 100k, 200k)...")
|
|
128
|
+
kmeans2 = MiniBatchKMeans(n_clusters=50000, random_state=42, batch_size=8192, max_iter=5, n_init="auto").fit(x)
|
|
129
|
+
kmeans3 = MiniBatchKMeans(n_clusters=100000, random_state=42, batch_size=8192, max_iter=10, n_init="auto").fit(x)
|
|
130
|
+
kmeans4 = MiniBatchKMeans(n_clusters=200000, random_state=42, batch_size=8192, max_iter=10, n_init="auto").fit(x)
|
|
131
|
+
|
|
132
|
+
data_dir = resources.files("dpst.data")
|
|
133
|
+
|
|
134
|
+
print(f"Saving cluster assets to: {data_dir}")
|
|
135
|
+
with open(data_dir / "50k.json", 'w') as f:
|
|
136
|
+
json.dump(kmeans2.cluster_centers_.tolist(), f)
|
|
137
|
+
with open(data_dir / "50k_counts.json", 'w') as f:
|
|
138
|
+
json.dump([Counter(kmeans2.labels_)[k] for k in sorted(Counter(kmeans2.labels_))], f)
|
|
139
|
+
|
|
140
|
+
with open(data_dir / "100k.json", 'w') as f:
|
|
141
|
+
json.dump(kmeans3.cluster_centers_.tolist(), f)
|
|
142
|
+
with open(data_dir / "100k_counts.json", 'w') as f:
|
|
143
|
+
json.dump([Counter(kmeans3.labels_)[k] for k in sorted(Counter(kmeans3.labels_))], f)
|
|
144
|
+
|
|
145
|
+
with open(data_dir / "200k.json", 'w') as f:
|
|
146
|
+
json.dump(kmeans4.cluster_centers_.tolist(), f)
|
|
147
|
+
with open(data_dir / "200k_counts.json", 'w') as f:
|
|
148
|
+
json.dump([Counter(kmeans4.labels_)[k] for k in sorted(Counter(kmeans4.labels_))], f)
|
|
149
|
+
|
|
150
|
+
print("Updating collection items with cluster allocations...")
|
|
151
|
+
for uid, b, c, d in tqdm(zip(all_ids, kmeans2.labels_, kmeans3.labels_, kmeans4.labels_), total=len(all_ids)):
|
|
152
|
+
props = {"fiftyk": int(b), "hundredk": int(c), "twohundredk": int(d)}
|
|
153
|
+
triples.data.update(uuid=uid, properties=props)
|
|
154
|
+
|
|
155
|
+
client.close()
|
|
156
|
+
print("Data preparation phase completely finished!")
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: dpst
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Differentially private document generation leveraging semantic triples.
|
|
5
|
+
Author-email: Stephen Meisenbacher <sjmeis@gtgd.com>
|
|
6
|
+
Maintainer-email: Stephen Meisenbacher <sjmeis@gtgd.com>
|
|
7
|
+
License: MIT License
|
|
8
|
+
|
|
9
|
+
Copyright (c) 2025-2026 Stephen Meisenbacher
|
|
10
|
+
|
|
11
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
12
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
13
|
+
in the Software without restriction, including without limitation the rights
|
|
14
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
15
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
16
|
+
furnished to do so, subject to the following conditions:
|
|
17
|
+
|
|
18
|
+
The above copyright notice and this permission notice shall be included in all
|
|
19
|
+
copies or substantial portions of the Software.
|
|
20
|
+
|
|
21
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
22
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
23
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
24
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
25
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
26
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
27
|
+
SOFTWARE.
|
|
28
|
+
Description-Content-Type: text/markdown
|
|
29
|
+
License-File: LICENSE
|
|
30
|
+
Requires-Dist: weaviate-client>=4.0.0
|
|
31
|
+
Requires-Dist: transformers
|
|
32
|
+
Requires-Dist: sentence-transformers
|
|
33
|
+
Requires-Dist: openie
|
|
34
|
+
Requires-Dist: tqdm
|
|
35
|
+
Requires-Dist: numpy
|
|
36
|
+
Requires-Dist: spacy
|
|
37
|
+
Requires-Dist: torch
|
|
38
|
+
Requires-Dist: datasketch
|
|
39
|
+
Requires-Dist: nltk
|
|
40
|
+
Provides-Extra: setup
|
|
41
|
+
Requires-Dist: datasets; extra == "setup"
|
|
42
|
+
Requires-Dist: scikit-learn; extra == "setup"
|
|
43
|
+
Dynamic: license-file
|
|
44
|
+
|
|
45
|
+
<div align="center">
|
|
46
|
+
|
|
47
|
+
# DP-ST
|
|
48
|
+
|
|
49
|
+
[](https://pypi.org/project/dpst/)
|
|
50
|
+
[](https://github.com/sjmeis/DPST/stargazers)
|
|
51
|
+
[](https://github.com/sjmeis/DPST/blob/main/LICENSE)
|
|
52
|
+
|
|
53
|
+
</div>
|
|
54
|
+
|
|
55
|
+
Code repository for the EMNLP 2025 paper: *Leveraging Semantic Triples for Private Document Generation with Local Differential Privacy Guarantees*
|
|
56
|
+
|
|
57
|
+
## Getting Started
|
|
58
|
+
### Installation
|
|
59
|
+
You can now install `DP-ST` directly as a Python package.
|
|
60
|
+
|
|
61
|
+
```bash
|
|
62
|
+
# Install the core runtime engine
|
|
63
|
+
pip install dpst
|
|
64
|
+
|
|
65
|
+
# Install the optional data-preparation extensions required for the initial setup
|
|
66
|
+
pip install "dpst[setup]"
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
**Note on Heavy Runtimes**: We highly recommend installing the correct flavor of PyTorch matching your hardware's CUDA capability before installing this package.
|
|
70
|
+
|
|
71
|
+
### Automated Database & Cluster Setup
|
|
72
|
+
In order to run `DP-ST`, you must first run the *preparation* stage as described in the paper. This includes booting up your local vector database, extracting triples from a public text corpus, clustering them, and storing them locally.
|
|
73
|
+
|
|
74
|
+
Ensure your local Weaviate instance is running, then execute the following automated command in your terminal:
|
|
75
|
+
|
|
76
|
+
```bash
|
|
77
|
+
dpst setup
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
**Note:** This can take a very long time! We recommend you set it and forget it. Alternatively, you can tweak the `max_rows` parameter of `initialize_database` to use less texts for the database preparation.
|
|
81
|
+
|
|
82
|
+
The above replaces the legacy workflow of manually executing `Triple2DB.ipynb` and `triple_cluster.ipynb`. The command will automatically stream the FineWeb public corpus dataset, extract/embed triples into Weaviate, run the MiniBatchKMeans allocations (50k, 100k, and 200k), and write the resulting cluster assets straight into the package data directory. If you prefer a more tailored approach, we still recommend using the notebooks.
|
|
83
|
+
|
|
84
|
+
## Usage
|
|
85
|
+
|
|
86
|
+
Running `DP-ST` is simple once the automated setup has completed:
|
|
87
|
+
|
|
88
|
+
```python
|
|
89
|
+
from dpst import DPST
|
|
90
|
+
|
|
91
|
+
# Initialize the engine (specify mode: "50k", "100k", or "200k")
|
|
92
|
+
X = DPST(mode="50k", model_checkpoint=MODEL_NAME, hf_token=TOKEN)
|
|
93
|
+
|
|
94
|
+
# Privatize your text corpus
|
|
95
|
+
private_texts = X.privatize([TEXTS], epsilon=DOC_PRIVACY_BUDGET)
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
`MODEL_NAME` refers to the model used for text reconstruction (i.e., the `Llama-3.2` models we use in the work), and `hf_token` is only necessary for gated models on Hugging Face.
|
|
99
|
+
|
|
100
|
+
## Running other DP Methods
|
|
101
|
+
In this repository (under the `comparison` directory), you will find a number of scripts (`*_perturb.py`) to reproduce the privatized texts as performed in our work.
|
|
102
|
+
|
|
103
|
+
The code for `DP-BART` and `DP-Prompt` can be found in `LLMDP.py`. `DP-MLM` can be found [here](https://github.com/sjmeis/DPMLM/) and the code for `TEM` can be found [here](https://github.com/sjmeis/MLDP/).
|
|
104
|
+
|
|
105
|
+
We also include the evaluation code for cosine similarity (`CS.py`) and G-Eval (`Geval.ipynb`), located in `evaluation`.
|
|
106
|
+
|
|
107
|
+
NOTE: in all provided notebooks, please make sure to include the correct libraries and link the paths accordingly. This is necessary for the code to run correctly!
|
|
108
|
+
|
|
109
|
+
## Citation
|
|
110
|
+
If you use this code in your research, please consider citing the published work:
|
|
111
|
+
|
|
112
|
+
```
|
|
113
|
+
@inproceedings{meisenbacher-etal-2025-leveraging,
|
|
114
|
+
title = "Leveraging Semantic Triples for Private Document Generation with Local Differential Privacy Guarantees",
|
|
115
|
+
author = "Meisenbacher, Stephen and
|
|
116
|
+
Chevli, Maulik and
|
|
117
|
+
Matthes, Florian",
|
|
118
|
+
editor = "Christodoulopoulos, Christos and
|
|
119
|
+
Chakraborty, Tanmoy and
|
|
120
|
+
Rose, Carolyn and
|
|
121
|
+
Peng, Violet",
|
|
122
|
+
booktitle = "Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing",
|
|
123
|
+
month = nov,
|
|
124
|
+
year = "2025",
|
|
125
|
+
address = "Suzhou, China",
|
|
126
|
+
publisher = "Association for Computational Linguistics",
|
|
127
|
+
url = "https://aclanthology.org/2025.emnlp-main.455/",
|
|
128
|
+
doi = "10.18653/v1/2025.emnlp-main.455",
|
|
129
|
+
pages = "8976--8992",
|
|
130
|
+
ISBN = "979-8-89176-332-6"
|
|
131
|
+
}
|
|
132
|
+
```
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
dpst/__init__.py,sha256=gK4XIDT6JcPOkYO5B4XXRloImX_byxGTdoiISGMQHlg,69
|
|
2
|
+
dpst/cli.py,sha256=FC5J4sDHSFw2cBZs_PGsmC-wcm_hHzzo0c-tlgZoWcg,912
|
|
3
|
+
dpst/core.py,sha256=MbmQnsb6N631UzM4qW2eSH1jokfUrCudD8nT6Rq4CUw,13599
|
|
4
|
+
dpst/setup_db.py,sha256=FkVMtQd8wNX2K2ehMmWpHkehEG4v_VH2RbtSP22NFOo,5846
|
|
5
|
+
dpst-0.1.0.dist-info/licenses/LICENSE,sha256=Fidi0t9GAtM6s1nRixU2cOHCfOgIsOJEkqNjgr7Zx8s,1101
|
|
6
|
+
dpst-0.1.0.dist-info/METADATA,sha256=2s8jD8E-KO4dv6qvC00fSN2F9P26qxNDMDGvJr5h8aA,6058
|
|
7
|
+
dpst-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
|
|
8
|
+
dpst-0.1.0.dist-info/entry_points.txt,sha256=NhFjOF2MhSoenGFK-kX8nK8aHzzoO4HJkCC3vwlgmzs,39
|
|
9
|
+
dpst-0.1.0.dist-info/top_level.txt,sha256=9NmKe0m8O75T_GUImjBprimRvxqi8YcJAurcCwr47JE,5
|
|
10
|
+
dpst-0.1.0.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025-2026 Stephen Meisenbacher
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
dpst
|