def analyze_sentences(
records:Iterable[Record],
text_field:str,
token_field:str="tokens",
entity_field:str="entities",
)->Iterable[Record]:
"""
Parses the text fields of all records using SciSpacy.
Requires that text_util:nlp and text_util:stopwords have both been loaded into
dask_process_global.
@param records: A partition of records to parse, each must contain `text_field`
@param text_field: The name of the field we wish to parse.
@param token_field: The output field for all basic tokens. These are
sub-records containing information such as POS tag and lemma.
@param entity_field: The output field for all entities, which are multi-token
phrases.
@return a list of records with token and entity fields
"""
nlp = dpg.get("text_util:nlp")
stopwords = dpg.get("text_util:stopwords")
res = []
for sent_rec, doc in zip(
records,
nlp.pipe(
map(
lambda x:x[text_field],
records
),
)
):
sent_rec[entity_field] = [
{
"tok_start": ent.start,
"tok_end": ent.end,
"cha_start": ent.start_char,
"cha_end": ent.end_char,
"label": ent.label_,
}
for ent in doc.ents
if ent.end - ent.start > 1 # don't want 1-gram ents
]
sent_rec[token_field] = [
{
"cha_start": tok.idx,
"cha_end": tok.idx + len(tok),
"lemma": tok.lemma_,
"pos": tok.pos_,
"tag": tok.tag_,
"dep": tok.dep_,
"stop": \
tok.lemma_ in stopwords or tok.text.strip().lower() in stopwords
}
for tok in doc
]
res.append(sent_rec)
return res
def set_index(
collection: str,
field_name: str,
index_type: MONGO_INDEX = pymongo.TEXT,
) -> None:
db = dpg.get("database:db")
db[collection].create_index([(field_name, index_type)])
def _init():
device = dpg.get("embedding_util:device")
tok = BertTokenizer.from_pretrained(bert_model)
model = BertModel.from_pretrained(bert_model)
model.eval()
model.to(device)
return (tok, model)
def _init():
device = dpg.get("embedding_util:device")
model = model_class(**model_kwargs)
model.load_state_dict(torch.load(
str(data_dir),
map_location=device,
))
model.eval()
return model
def get(values: Iterable[str],
collection: str,
field_name: str,
desired_fields: List[str] = None,
**kwargs) -> Iterable[Record]:
db = dpg.get("database:db")
if desired_fields is None:
proj_obj = {}
else:
proj_obj = {field_name: 1 for field_name in desired_fields}
proj_obj["_id"] = 0
return [db[collection].find_one({field_name: v}, proj_obj) for v in values]
def put(records: Iterable[Record], collection: str, **kwargs) -> None:
"""
Inserts all the records, returns the resulting count. Note that
additional_args are allowed to specify delayed dependencies.
"""
db = dpg.get("database:db")
for r in records:
try:
db[collection].insert(r)
except pymongo.errors.InvalidOperation as e:
print("Encountered non-fatal issue:", e)
pass
def apply_sentence_classifier_to_part(
records: Iterable[Record],
batch_size: int,
sentence_classifier_name="sentence_classifier",
predicted_type_suffix=":pred",
sentence_type_field="sent_type",
) -> Iterable[Record]:
device = dpg.get("embedding_util:device")
model = dpg.get(f"embedding_util:{sentence_classifier_name}")
res = []
for rec_batch in iter_to_batches(records, batch_size):
model_input = torch.stack([
record_to_sentence_classifier_input(r) for r in rec_batch
]).to(device)
predicted_labels = sentence_classifier_output_to_labels(
model(model_input))
for r, lbl in zip(rec_batch, predicted_labels):
r[sentence_type_field] = lbl + predicted_type_suffix
res.append(r)
print(len(res))
return res
def embed_records(
records: Iterable[Record],
batch_size: int,
text_field: str,
max_sequence_length: int,
out_embedding_field: str = "embedding",
show_pbar: bool = False,
) -> Iterable[Record]:
"""
Introduces an embedding field to each record, indicated the bert embedding
of the supplied text field.
"""
dev = dpg.get("embedding_util:device")
tok, model = dpg.get("embedding_util:tok,model")
res = []
# PBar is necessary when using embed_records in helper scripts.
pbar = tqdm(
iter_to_batches(records, batch_size),
total=int(len(records) / batch_size),
disable=not show_pbar,
)
for batch in pbar:
texts = list(map(lambda x: x[text_field], batch))
sequs = pad_sequence(
sequences=[
torch.tensor(tok.encode(t)[:max_sequence_length])
for t in texts
],
batch_first=True,
).to(dev)
with torch.no_grad():
embs = (model(sequs)[-2].mean(axis=1).cpu().detach().numpy())
for record, emb in zip(batch, embs):
record[out_embedding_field] = emb
res.append(record)
return res
def apply_faiss_to_edges(
hash_and_embedding:Iterable[Record],
)->Iterable[nx.Graph]:
# The only reason we need parts_written_to_db is to make sure that the
# writing happens before this point
index = dpg.get(f"knn_util:faiss_{faiss_index_name}")
inverted_index = {}
graph = nx.Graph()
for batch in iter_to_batches(hash_and_embedding, batch_size):
hashes, embeddings = records_to_ids_and_embeddings(
records=batch,
)
_, neighs_per_root = index.search(embeddings, num_neighbors)
hashes = hashes.tolist() + flatten_list(neighs_per_root.tolist())
hashes = list(set(hashes) - set(inverted_index.keys()))
graph_keys = database_util.get(
values=hashes,
collection=inverted_index_collection,
field_name="hash",
desired_fields=["strid"]
)
for k, v in zip(hashes, graph_keys):
inverted_index[k] = v["strid"]
# Create records
for root_idx, neigh_indices in zip(hashes, neighs_per_root):
root = inverted_index[root_idx]
if root is None:
continue
for neigh_idx in neigh_indices:
if neigh_idx == root_idx:
continue
neigh = inverted_index[neigh_idx]
if neigh is None:
continue
graph.add_edge(root, neigh, weight=weight)
graph.add_edge(neigh, root, weight=weight)
return [graph]
def clear_collection(collection: str, **kwargs) -> None:
db = dpg.get("database:db")
db[collection].drop()