Kf/eschatology processing

.pre-commit-config.yaml

@@ -21,3 +21,9 @@ repos:
     rev: v0.5.7
     hooks:
       - id: ruff
+
+  - repo: https://github.com/pre-commit/mirrors-prettier
+    rev: v3.1.0
+    hooks:
+      - id: prettier
+        files: "visualizer/.*"

config/eschatology.toml

@@ -2,7 +2,7 @@
 language = "en"
 
 [preprocessing]
-enabled = true
+enabled = false
 doc_type = "csv"
 
 [preprocessing.extra]
@@ -11,9 +11,14 @@ text_column = "body"
 timestamp_column = "timestamp"
 
 [docprocessing]
-enabled = true
-batch_size = 5
-prefer_gpu_for_coref = false
+enabled = false
+batch_size = 50
+prefer_gpu_for_coref = true
+n_process = 1
 
 [corpusprocessing]
-enabled = true
+enabled = false
+
+[databasepopulation]
+enabled = true
+clear_and_write = true

config/template.toml

@@ -12,13 +12,14 @@ metadata_fields = ["*"]
 
 [preprocessing.extra]
 # specific extra arguments for your preprocessor, e.g. context length for tweets or
-# or field specification for CSVs
+# field specification for CSVs
 
 [docprocessing]
 enabled = true
 batch_size = 25
 continue_from_last = true
 triplet_extraction_method = "multi2oie/prompting"
+n_process = 1 # can be set to 2 or more for multiprocess ofloading to GPU; otherwise might not make sense
 
 [corpusprocessing]
 enabled = true
@@ -27,6 +28,6 @@ dimensions = 100 # leave out to skip dimensionality reduction
 n_neighbors = 15 # used for dimensionality reduction
 
 [corpusprocessing.thresholds]  # leave out for automatic estimation
-min_cluster_size = 3  # unused if auto_thresholds is true
-min_samples = 3  # unused if auto_thresholds is true
-min_topic_size = 5  # unused if auto_thresholds is true
+min_label_occurrence = 3
+min_cluster_size = 3
+min_samples = 3

docs/tutorials/overview.ipynb

@@ -71,7 +71,7 @@
     "    assert isinstance(sent._.coref_clusters[0], tuple)\n",
     "    assert isinstance(sent._.coref_clusters[0][0], int)\n",
     "    assert isinstance(sent._.coref_clusters[0][1], Span)\n",
-    "    sent._.resolve_coref  # get resolved coref"
+    "    sent._.resolved_text  # get resolved coref"
    ]
   },
   {

paper/extract_triplets_newspapers.py

@@ -88,7 +88,7 @@ def process_file(
 
         # Resolve coreference
         coref_docs = nlp_coref.pipe(normalized_article)
-        resolved_docs = (d._.resolve_coref for d in coref_docs)
+        resolved_docs = (d._.resolved_text for d in coref_docs)
 
         # Extract relations
         docs = nlp.pipe(resolved_docs)

paper/extract_triplets_tweets.py

@@ -99,7 +99,7 @@ def concat_resolve_unconcat_contexts(file_path: str):
 
     coref_nlp = build_coref_pipeline()
     coref_docs = coref_nlp.pipe(context_tweets)
-    resolved_docs = (d._.resolve_coref for d in coref_docs)
+    resolved_docs = (d._.resolved_text for d in coref_docs)
 
     resolved_tweets = (tweet_from_context_text(tweet) for tweet in resolved_docs)
     return resolved_tweets
@@ -240,7 +240,7 @@ def prompt_gpt3(
     for i, batch in enumerate(batch_generator(concatenated_tweets, batch_size)):
         start = time.time()
         coref_docs = coref_nlp.pipe(batch)
-        resolved_docs = (d._.resolve_coref for d in coref_docs)
+        resolved_docs = (d._.resolved_text for d in coref_docs)
         resolved_target_tweets = (
             tweet_from_context_text(tweet) for tweet in resolved_docs
         )

pyproject.toml

@@ -45,7 +45,9 @@ dependencies = [
     "sentence-transformers",
     "stop-words",
     "bs4",
-    "toml"
+    "toml",
+    "fastcoref",
+    "sqlalchemy"
 ]
 
 [project.license]
@@ -94,6 +96,7 @@ content-type = "text/markdown"
 "prompt_relation_extraction" = "conspiracies.docprocessing.relationextraction.gptprompting:create_prompt_relation_extraction_component"
 "relation_extractor" = "conspiracies.docprocessing.relationextraction.multi2oie:make_relation_extractor"
 "allennlp_coref" = "conspiracies.docprocessing.coref:create_coref_component"
+"safe_fastcoref" = "conspiracies.docprocessing.coref.safefastcoref:create_safe_fastcoref"
 "heads_extraction" = "conspiracies.docprocessing.headwordextraction:create_headwords_component"
 
 

src/conspiracies/corpusprocessing/clustering.py

@@ -1,12 +1,15 @@
-from collections import defaultdict
-from typing import List, Callable, Any, Hashable, Dict
+import math
+import os
+from collections import defaultdict, Counter
+from pathlib import Path
+from typing import List, Callable, Any, Hashable, Dict, Union
 
 import networkx
 import numpy as np
 from hdbscan import HDBSCAN
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
-from sklearn.preprocessing import StandardScaler
+from tqdm import tqdm
 from umap import UMAP
 
 from conspiracies.common.modelchoice import ModelChoice
@@ -45,13 +48,17 @@ def __init__(
         min_cluster_size: int = 5,
         min_samples: int = 3,
         embedding_model: str = None,
+        cache_location: Path = None,
     ):
         self.language = language
         self.n_dimensions = n_dimensions
         self.n_neighbors = n_neighbors
         self.min_cluster_size = min_cluster_size
         self.min_samples = min_samples
         self._embedding_model = embedding_model
+        self.cache_location = cache_location
+        if self.cache_location is not None:
+            os.makedirs(self.cache_location, exist_ok=True)
 
     def _get_embedding_model(self):
         # figure out embedding model if not given explicitly
@@ -97,53 +104,88 @@ def _combine_clusters(
 
         return merged_clusters
 
-    def _cluster(
+    def _cluster_via_embeddings(
         self,
-        fields: List[TripletField],
+        labels: List[str],
+        cache_name: str = None,
+        show_progress: bool = True,
     ):
-        model = self._get_embedding_model()
-        print("Creating embeddings:")
-        embeddings = model.encode(
-            [field.text for field in fields],
-            show_progress_bar=True,
+        emb_cache = (
+            Path(self.cache_location, f"embeddings-{cache_name}.npy")
+            if self.cache_location and cache_name
+            else None
         )
-        embeddings = StandardScaler().fit_transform(embeddings)
+        if emb_cache and emb_cache.exists():
+            print(
+                "Reusing cached embeddings! Delete cache if this is not supposed to happen.",
+            )
+            embeddings = np.load(emb_cache)
+        else:
+            model = self._get_embedding_model()
+
+            counter = Counter((field for field in labels))
+            condensed = [
+                field
+                for field, count in counter.items()
+                for _ in range(math.ceil(count / 1000))
+            ]
+            embeddings = model.encode(
+                condensed,
+                normalize_embeddings=True,
+                show_progress_bar=show_progress,
+            )
+            if emb_cache:
+                np.save(emb_cache, embeddings)
 
         if self.n_dimensions is not None:
-            print("Reducing embedding space")
-            reducer = UMAP(n_components=self.n_dimensions, n_neighbors=self.n_neighbors)
-            embeddings = reducer.fit_transform(embeddings)
+            reduced_emb_cache = (
+                Path(
+                    self.cache_location,
+                    f"embeddings-{cache_name}-red{self.n_dimensions}.npy",
+                )
+                if self.cache_location and cache_name
+                else None
+            )
+            if reduced_emb_cache and reduced_emb_cache.exists():
+                print(
+                    "Reusing cached reduced embeddings! Delete cache if this is not supposed to happen.",
+                )
+                embeddings = np.load(reduced_emb_cache)
+            else:
+                print("Reducing embedding space ...")
+                reducer = UMAP(
+                    n_components=self.n_dimensions,
+                    n_neighbors=self.n_neighbors,
+                )
+                embeddings = reducer.fit_transform(embeddings)
+                if self.cache_location:
+                    np.save(reduced_emb_cache, embeddings)
 
-        print("Clustering ...")
         hdbscan_model = HDBSCAN(
             min_cluster_size=self.min_cluster_size,
+            max_cluster_size=self.min_cluster_size
+            * 10,  # somewhat arbitrary, mostly to avoid mega clusters that suck up everything
             min_samples=self.min_samples,
         )
         hdbscan_model.fit(embeddings)
 
         clusters = defaultdict(list)
         for field, embedding, label, probability in zip(
-            fields,
+            labels,
             embeddings,
             hdbscan_model.labels_,
             hdbscan_model.probabilities_,
         ):
             # skip noise and low confidence
-            if label == -1 or probability < 0.1:
+            if label == -1 or probability < 0.5:
                 continue
             clusters[label].append((field, embedding))
 
         merged = self._combine_clusters(
             list(clusters.values()),
-            get_combine_key=lambda t: t[0].text,
+            get_combine_key=lambda t: t[0],
         )
 
-        # too risky with false positives from this
-        # merged = self._combine_clusters(
-        #     merged,
-        #     get_combine_key=lambda t: t[0].head,
-        # )
-
         # sort by how "prototypical" a member is in the cluster
         for cluster in merged:
             mean = np.mean(np.stack([t[1] for t in cluster]), axis=0)
@@ -153,11 +195,69 @@ def _cluster(
         return [[t[0] for t in cluster] for cluster in merged]
 
     @staticmethod
-    def _mapping_to_first_member(clusters: List[List[TripletField]]) -> Dict[str, str]:
+    def _cluster_via_normalization(
+        labels: List[str],
+        top: Union[int, float] = 1.0,
+        restrictive_labels=True,
+    ) -> List[List[str]]:
+        counter = Counter((label for label in labels))
+        if isinstance(top, float):
+            top = int(top * len(counter))
+
+        norm_map = {
+            label: " "
+            + label.lower()
+            + " "  # surrounding spaces avoids matches like evil <-> devil
+            for label in counter.keys()
+        }
+        cluster_map = {
+            label: []
+            for label, count in counter.most_common(top)
+            # FIXME: hack due to lack of NER and lemmas at the time of writing
+            if not restrictive_labels
+            or len(label) >= 4
+            and label[0].isupper()
+            or len(label.split()) > 1
+        }
+
+        for label in counter.keys():
+            norm_label = norm_map[label]
+            matches = [
+                substring
+                for substring in cluster_map.keys()
+                if norm_map[substring] in norm_label
+            ]
+            if not matches:
+                continue
+
+            best_match = min(
+                matches,
+                key=lambda substring: len(norm_map[substring]),
+            )
+            if best_match != label:
+                cluster_map[best_match].append(label)
+
+        clusters = [
+            [main_label] + alt_labels
+            for main_label, alt_labels in cluster_map.items()
+            if alt_labels
+        ]
+        return clusters
+
+    @staticmethod
+    def _mapping_to_first_member(
+        clusters: List[List[Union[TripletField, str]]],
+    ) -> Dict[str, str]:
+        def get_text(member: Union[TripletField, str]):
+            if isinstance(member, TripletField):
+                return member.text
+            else:
+                return member
+
         return {
-            member: cluster[0].text
+            member: get_text(cluster[0])
             for cluster in clusters
-            for member in set(member.text for member in cluster)
+            for member in set(get_text(member) for member in cluster)
         }
 
     def create_mappings(self, triplets: List[Triplet]) -> Mappings:
@@ -166,10 +266,33 @@ def create_mappings(self, triplets: List[Triplet]) -> Mappings:
         entities = subjects + objects
         predicates = [triplet.predicate for triplet in triplets]
 
+        # FIXME: clustering gets way to aggressive for many triplets
+        # print("Creating mappings for entities")
+        # entity_clusters = self._cluster(entities, "entities")
+        # print("Creating mappings for predicates")
+        # predicate_clusters = self._cluster(predicates, "predicates")
+
         print("Creating mappings for entities")
-        entity_clusters = self._cluster(entities)
+        entity_clusters = self._cluster_via_normalization(
+            [e.text for e in entities],
+            0.2,
+        )
+        entity_clusters = [
+            sub_cluster
+            for cluster in tqdm(entity_clusters, desc="Creating sub-clusters")
+            for sub_cluster in (
+                self._cluster_via_embeddings(cluster, show_progress=False)
+                if len(cluster) > 10
+                else [cluster]
+            )
+        ]
+
         print("Creating mappings for predicates")
-        predicate_clusters = self._cluster(predicates)
+        predicate_clusters = self._cluster_via_normalization(
+            [p.text for p in predicates],
+            top=0.2,
+            restrictive_labels=False,
+        )
 
         mappings = Mappings(
             entities=self._mapping_to_first_member(entity_clusters),