Upload model

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "SimilarityModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration.SimilarityModelConfig",
+    "AutoModel": "similarity_model.SimilarityModel"
+  },
+  "device": "cpu",
+  "embedding_model_config": {
+    "add_upper": true,
+    "d": 128,
+    "device": "cpu",
+    "model_class": "roberta",
+    "model_name": "roberta-base",
+    "normalize": true,
+    "pooling": "pooler",
+    "prompt": "",
+    "upper_case": false
+  },
+  "model_type": "roberta",
+  "score_model_config": {
+    "alpha": 50
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.27.4",
+  "weighting_function_config": {
+    "weighting_exponent": 0.5
+  }
+}

configuration.py

@@ -0,0 +1,30 @@
+from transformers import PretrainedConfig
+
+
+class SimilarityModelConfig(PretrainedConfig):
+    model_type = 'roberta'
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        
+        self.embedding_model_config = kwargs.get("embedding_model_config")
+        self.score_model_config = kwargs.get("score_model_config")
+        self.weighting_function_config = kwargs.get("weighting_function_config")
+
+
+nama_base = SimilarityModelConfig(
+    embedding_model_config={
+        "model_class": 'roberta',
+        "model_name":'roberta-base',
+        "pooling": 'pooler',
+        "normalize":True,
+        "d":128,
+        "prompt":'',
+        "device":'cpu',
+        "add_upper": True,
+        "upper_case":False
+        },
+    score_model_config={"alpha": 50},
+    weighting_function_config={"weighting_exponent": 0.5},
+    device="cpu",
+)
+

embedding_model.py

@@ -0,0 +1,121 @@
+import torch
+from transformers import AutoTokenizer, RobertaModel
+
+class EmbeddingModel(torch.nn.Module):
+
+
+    tokenizers = {'roberta': RobertaModel}
+    """
+    A basic wrapper around a Hugging Face transformer model.
+    Takes a string as input and produces an embedding vector of size d.
+    """
+    def __init__(self, config, **kwargs):
+
+        super().__init__()
+
+        self.model_class = self.tokenizers.get(config.get("model_class").lower())
+        self.model_name = config.get("model_name")
+        self.pooling = config.get("pooling")
+        self.normalize = config.get("normalize")
+        self.d = config.get("d")
+        self.prompt = config.get("prompt")
+        self.add_upper = config.get("add_upper")
+        self.upper_case = config.get("upper_case")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+
+        try:
+            self.transformer = self.model_class.from_pretrained(self.model_name)
+        except OSError:
+            self.transformer = self.model_class.from_pretrained(self.model_name,from_tf=True)
+
+        self.dropout = torch.nn.Dropout(0.5)
+
+        if self.d:
+            # Project embedding to a lower dimension
+            # Initialization based on random projection LSH (preserves approximate cosine distances)
+            self.projection = torch.nn.Linear(self.transformer.config.hidden_size,self.d)
+            torch.nn.init.normal_(self.projection.weight)
+            torch.nn.init.constant_(self.projection.bias,0)
+
+        self.to(config.get("device"))
+
+    def to(self,device):
+        super().to(device)
+        self.device = device
+
+    def encode(self,strings):
+        if self.prompt is not None:
+            strings = [self.prompt + s for s in strings]
+        if self.add_upper:
+            strings = [s + ' </s> ' + s.upper() for s in strings]
+        if self.upper_case:
+            strings = [s + ' </s> ' + s.upper() for s in strings]
+
+        try:
+            encoded = self.tokenizer(strings,padding=True,truncation=True)
+        except Exception as e:
+            print(strings)
+            raise Exception(e)
+        input_ids = torch.tensor(encoded['input_ids']).long()
+        attention_mask = torch.tensor(encoded['attention_mask'])
+
+        return input_ids,attention_mask
+
+    def forward(self,strings):
+
+        with torch.no_grad():
+            input_ids,attention_mask = self.encode(strings)
+
+            input_ids = input_ids.to(device=self.device)
+            attention_mask = attention_mask.to(device=self.device)
+
+        # with amp.autocast(self.amp):
+        batch_out = self.transformer(input_ids=input_ids,
+                                        attention_mask=attention_mask,
+                                        return_dict=True)
+
+        if self.pooling == 'pooler':
+            v = batch_out['pooler_output']
+        elif self.pooling == 'mean':
+            h = batch_out['last_hidden_state']
+
+            # Compute mean of unmasked token vectors
+            h = h*attention_mask[:,:,None]
+            v = h.sum(dim=1)/attention_mask.sum(dim=1)[:,None]
+
+        if self.d:
+            v = self.projection(v)
+
+        if self.normalize:
+            v = v/torch.sqrt((v**2).sum(dim=1)[:,None])
+
+        return v
+
+    def config_optimizer(self,transformer_lr=1e-5,projection_lr=1e-4):
+
+        parameters = list(self.named_parameters())
+        grouped_parameters = [
+                {
+                    'params': [param for name,param in parameters if name.startswith('transformer') and name.endswith('bias')],
+                    'weight_decay_rate': 0.0,
+                    'lr':transformer_lr,
+                    },
+                {
+                    'params': [param for name,param in parameters if name.startswith('transformer') and not name.endswith('bias')],
+                    'weight_decay_rate': 0.0,
+                    'lr':transformer_lr,
+                    },
+                {
+                    'params': [param for name,param in parameters if name.startswith('projection')],
+                    'weight_decay_rate': 0.0,
+                    'lr':projection_lr,
+                    },
+                ]
+
+        # Drop groups with lr of 0
+        grouped_parameters = [p for p in grouped_parameters if p['lr']]
+
+        optimizer = torch.optim.AdamW(grouped_parameters)
+
+        return optimizer

embeddings.py

@@ -0,0 +1,637 @@
+import pandas as pd
+import numpy as np
+from tqdm import tqdm
+from copy import copy
+from collections import Counter
+import torch
+from zipfile import ZipFile
+import pickle
+from io import BytesIO
+
+from .match_groups import MatchGroups
+
+
+class Embeddings(torch.nn.Module):
+    """
+    Stores embeddings for a fixed array of strings and provides methods for
+    clustering the strings to create MatchGroups objects according to different
+    algorithms.
+    """
+    def __init__(self,strings,V,score_model,weighting_function,counts,device='cpu'):
+        super().__init__()
+
+        self.strings = np.array(list(strings))
+        self.string_map = {s:i for i,s in enumerate(strings)}
+        self.V = V
+        self.counts = counts
+        self.w = weighting_function(counts)
+        self.score_model = score_model
+        self.weighting_function = weighting_function
+        self.device = device
+
+        self.to(device)
+
+    def __repr__(self):
+        return f'<nama.Embeddings containing {self.V.shape[1]}-d vectors for {len(self)} strings'
+
+    def to(self,device):
+        super().to(device)
+        self.V = self.V.to(device)
+        self.counts = self.counts.to(device)
+        self.w = self.w.to(device)
+        self.score_model.to(device)
+        self.device = device
+
+    def save(self,f):
+        """
+        Save embeddings in a simple custom zipped archive format (torch.save
+        works too, but it requires huge amounts of memory to serialize large
+        embeddings objects).
+        """
+        with ZipFile(f,'w') as zip:
+
+            # Write score model
+            zip.writestr('score_model.pkl',pickle.dumps(self.score_model))
+
+            # Write score model
+            zip.writestr('weighting_function.pkl',pickle.dumps(self.weighting_function))
+
+            # Write string info
+            strings_df = pd.DataFrame().assign(
+                                        string=self.strings,
+                                        count=self.counts.to('cpu').numpy())
+            zip.writestr('strings.csv',strings_df.to_csv(index=False))
+
+            # Write embedding vectors
+            byte_io = BytesIO()
+            np.save(byte_io,self.V.to('cpu').numpy(),allow_pickle=False)
+            zip.writestr('V.npy',byte_io.getvalue())
+
+    def __getitem__(self,arg):
+        """
+        Slice a Match Groups object
+        """
+        if isinstance(arg,slice):
+            i = arg
+        elif isinstance(arg, MatchGroups):
+            return self[arg.strings()]
+        elif hasattr(arg,'__iter__'):
+            # Return a subset of the embeddings and their weights
+            string_map = self.string_map
+            i = [string_map[s] for s in arg]
+
+            if i == list(range(len(self))):
+                # Just selecting the whole match groups object - no need to slice the embedding
+                return copy(self)
+        else:
+            raise ValueError(f'Unknown slice input type ({type(input)}). Can only slice Embedding with a slice, match group, or iterable.')
+
+        new = copy(self)
+        new.strings = self.strings[i]
+        new.V = self.V[i]
+        new.counts = self.counts[i]
+        new.w = self.w[i]
+        new.string_map = {s:i for i,s in enumerate(new.strings)}
+
+        return new
+
+    def embed(self,grouping):
+        """
+        Construct updated Embeddings with counts from the input MatchGroups
+        """
+        new = self[grouping]
+        new.counts = torch.tensor([grouping.counts[s] for s in new.strings],device=self.device)
+        new.w = new.weighting_function(new.counts)
+
+        return new
+
+    def __len__(self):
+        return len(self.strings)
+
+    def _group_to_ids(self,grouping):
+        group_id_map = {g:i for i,g in enumerate(grouping.groups.keys())}
+        group_ids = torch.tensor([group_id_map[grouping[s]] for s in self.strings]).to(self.device)
+        return group_ids
+
+    def _ids_to_group(self,group_ids):
+        if isinstance(group_ids,torch.Tensor):
+            group_ids = group_ids.to('cpu').numpy()
+
+        strings = self.strings
+        counts = self.counts.to('cpu').numpy()
+
+        # Sort by group and string count
+        g_sort = np.lexsort((counts,group_ids))
+        group_ids = group_ids[g_sort]
+        strings = strings[g_sort]
+        counts = counts[g_sort]
+
+        # Identify group boundaries and split locations
+        split_locs = np.nonzero(group_ids[1:] != group_ids[:-1])[0] + 1
+
+        # Get grouped strings as separate arrays
+        groups = np.split(strings,split_locs)
+
+        # Build the groupings
+        grouping = MatchGroups()
+        grouping.counts = Counter({s:int(c) for s,c in zip(strings,counts)})
+        grouping.labels = {s:g[-1] for g in groups for s in g}
+        grouping.groups = {g[-1]:list(g) for g in groups}
+
+        return grouping
+
+    @torch.no_grad()
+    def _fast_unite_similar(self,group_ids,threshold=0.5,progress_bar=True,batch_size=64):
+
+        V = self.V
+        cos_threshold = self.score_model.score_to_cos(threshold)
+
+        for batch_start in tqdm(range(0,len(self),batch_size),
+                                    delay=1,desc='Predicting matches',disable=not progress_bar):
+
+            i_slice = slice(batch_start,batch_start+batch_size)
+            j_slice = slice(batch_start+1,None)
+
+            g_i = group_ids[i_slice]
+            g_j = group_ids[j_slice]
+
+            # Find j's with jaccard > threshold ("matches")
+            batch_matched = (V[i_slice]@V[j_slice].T >= cos_threshold) \
+                            * (g_i[:,None] != g_j[None,:])
+
+            for k,matched in enumerate(batch_matched):
+                if matched.any():
+                    # Get the group ids of the matched j's
+                    matched_groups = g_j[matched]
+
+                    # Identify all embeddings in these groups
+                    ids_to_group = torch.isin(group_ids,matched_groups)
+
+                    # Assign all matched embeddings to the same group
+                    group_ids[ids_to_group] = g_i[k].clone()
+
+        return self._ids_to_group(group_ids)
+
+    @torch.no_grad()
+    def unite_similar(self,
+                threshold=0.5,
+                group_threshold=None,
+                always_match=None,
+                never_match=None,
+                batch_size=64,
+                progress_bar=True,
+                always_never_conflicts='warn',
+                return_united=False):
+
+        """
+        Unite embedding strings according to predicted pairwise similarity.
+
+        - "theshold" sets the minimimum match similarity required to unite two strings.
+            - Note that strings with similarity<threshold can end up matched if they are
+              linked by a chain of sufficiently similar strings (matching is transitive).
+              "group_threshold" can be used to add an additional constraing on the minimum
+              similarity within each group.
+        - "group_threshold" sets the minimum similarity required within a single group.
+        - "always_match" takes any argument that can be used to unite strings. These 
+            strings will always be matched.
+        - "never_match" takes a set, or a list of sets, where each set indicates two or
+            more strings that should never be united with each other (these strings may 
+            still be united with other strings).
+        - "always_never_conflicts" determines how to handle conflicts between 
+            "always_match" and "never_match":
+            - always_never_conflicts="warn": Check for conflicts and print a warning
+                if any are found (default)
+            - always_never_conflicts="raise": Check for conflicts and raise an error
+                if any are found
+            - always_never_conflicts="ignore": Do not check for conflicts ("always_match"
+              will take precedence)
+
+        If "group_threshold" or "never_match" arguments are supplied, strings pairs are
+        united in order of similarity. Highest similarity strings are matched first, and 
+        before each time a new pair of strings is united, the function checks if this will
+        result in grouping any two strings with similarity<group_threshold. If so, this
+        pair is skipped. This version of the algorithm requires more memory and processing
+        time, but guaruntees deterministic output that is consistent with the constraints.
+            
+        returns: MatchGroups object
+        """
+        if group_threshold and group_threshold < threshold:
+            raise ValueError('group_threshold must be greater than or equal to threshold')
+
+        group_ids = torch.arange(len(self)).to(self.device)
+        
+        if always_match is not None:
+            always_grouping = (MatchGroups(self.strings)
+                            .unite(always_match))
+            always_match_labels = always_grouping.labels
+
+
+        # Use a simpler, faster prediction algorithm if possible
+        if not (return_united or group_threshold or (never_match is not None)):
+            if always_match is not None:
+                group_ids = self._group_to_ids(always_grouping)
+
+            return self._fast_unite_similar(
+                        group_ids=group_ids,
+                        threshold=threshold,
+                        batch_size=batch_size,
+                        progress_bar=progress_bar)
+
+        if never_match is not None:
+            # Ensure never_match is a nested list
+            if all(isinstance(s,str) for s in never_match):
+                never_match = [never_match]
+
+            if always_match is not None:
+
+                assert always_never_conflicts in ['raise','warn','ignore']
+                
+                if always_never_conflicts != 'ignore':
+
+                    # Find conflicts between never_match and always_match groups
+                    conflicts = []
+                    for i,g in enumerate(never_match):
+                        g = sorted(list(g))
+                        g_labels = [always_match_labels.get(s,s) for s in g]
+                        if len(set(g_labels)) < len(g):
+                            df = (pd.DataFrame()
+                                  .assign(
+                                    string=g,
+                                    never_match_group=i,
+                                    always_match_group=g_labels
+                                    ))
+                            conflicts.append(df)
+
+                    if conflicts:
+                        conflicts_df = pd.concat(conflicts)
+
+                        if always_never_conflicts == 'warn':
+                            print(f'Warning: The following never_match groups are in conflict with always_match groups:\n{conflicts_df}')
+                            print('Conflicted never_match relationships will be ignored')
+                        else:
+                            raise ValueError(f'The following never_match groups are in conflict with always_match groups\n{conflicts_df}')
+                                
+
+                # If always_match, collapse to group labels that should not match
+                # Note: Implicitly letting always_match over-ride never_match here
+                never_match = [{always_match_labels[s] for s in g if s in always_match_labels} for g in never_match]
+            
+            else:
+                # Otherwise just use the strings themselves as labels
+                never_match = [set(s) for s in never_match]
+
+        # Convert thresholds from scores to raw cosine distances
+        V = self.V
+        cos_threshold = self.score_model.score_to_cos(threshold)
+        if group_threshold is not None:
+            separate_cos = self.score_model.score_to_cos(group_threshold)
+
+        # First collect all pairs to match (can be memory intensive!)
+        matches = []
+        cos_scores = []
+        for batch_start in tqdm(range(0,len(self),batch_size),
+                                    desc='Scoring pairs',
+                                    delay=1,disable=not progress_bar):
+
+            i_slice = slice(batch_start,batch_start+batch_size)
+            j_slice = slice(batch_start+1,None)
+
+            # Find j's with jaccard > threshold ("matches")
+            batch_cos = V[i_slice]@V[j_slice].T
+
+            # Search upper diagonal entries only
+            # (note j_slice starting index is offset by one)
+            batch_cos = torch.triu(batch_cos)
+
+            bi,bj = torch.nonzero(batch_cos >= cos_threshold,as_tuple=True)
+
+            if len(bi):
+                # Convert batch index locations to global index locations
+                i = bi + batch_start
+                j = bj + batch_start + 1
+
+                cos = batch_cos[bi,bj]
+
+                # Can skip strings that are already matched in the base grouping
+                unmatched = group_ids[i] != group_ids[j]
+                i = i[unmatched]
+                j = j[unmatched]
+                cos = cos[unmatched]
+
+                if len(i):
+                    batch_matches = torch.hstack([i[:,None],j[:,None]])
+
+                    matches.append(batch_matches.to('cpu').numpy())
+                    cos_scores.append(cos.to('cpu').numpy())
+
+        # Unite potential match pairs in priority order, while respecting
+        # the group_threshold and never_match arguments
+        united = []
+        if matches:
+            matches = np.vstack(matches)
+            cos_scores = np.hstack(cos_scores).T
+
+            # Sort matches in descending order of score
+            m_sort = cos_scores.argsort()[::-1]
+            matches = matches[m_sort]
+
+            if return_united:
+                # Save cos scores for later return
+                cos_scores_df = pd.DataFrame(matches,columns=['i','j'])
+                cos_scores_df['cos'] = cos_scores[m_sort]
+
+            # Set up tensors
+            matches = torch.tensor(matches).to(self.device)
+            
+            # Set-up per-string tracking of never-match relationships
+            if never_match is not None:
+                never_match_map = {s:sep for sep in never_match for s in sep}
+                
+                if always_match is not None:
+                    # If always_match, we use group labels instead of the strings themselves
+                    never_match_array = np.array([never_match_map.get(always_match_labels[s],set()) for s in self.strings])
+                else:
+                    never_match_array = np.array([never_match_map.get(s,set()) for s in self.strings])                   
+            
+
+            n_matches = matches.shape[0]
+            with tqdm(total=n_matches,desc='Uniting matches',
+                        delay=1,disable=not progress_bar) as p_bar:
+
+                while len(matches):
+
+                    # Select the current match pair and remove it from the queue
+                    match_pair = matches[0]
+                    matches = matches[1:]
+
+                    # Get the groups of the current match pair
+                    g = group_ids[match_pair]
+                    g0 = group_ids == g[0]
+                    g1 = group_ids == g[1]
+
+                    # Identify which strings should be united
+                    to_unite = g0 | g1
+
+                    # Flag whether the new group will have three or more strings
+                    singletons = to_unite.sum() < 3
+
+                    # Start by asuming that we can match this pair
+                    unite_ok = True
+
+                    # Check whether uniting this pair will unite any never_match strings/labels
+                    if never_match is not None:
+                        never_0 = never_match_array[match_pair[0]]
+                        never_1 = never_match_array[match_pair[1]]
+
+                        if never_0 and never_1 and (never_0 & never_1):
+                            # Here we make use of the fact that any pair of never_match strings/labels
+                            # will appear in both never_0 and never_1 if one string/label is in each group
+                            unite_ok = False
+
+                    # Check whether the uniting the pair will violate the group_threshold
+                    # (impossible if the strings are singletons)
+                    if unite_ok and group_threshold and not singletons:
+                        V0 = V[g0,:]
+                        V1 = V[g1,:]
+
+                        unite_ok = (V0@V1.T).min() >= separate_cos
+
+
+                    if unite_ok:
+
+                        # Unite groups
+                        group_ids[to_unite] = g[0]
+
+                        if never_match and (never_0 or never_1):
+                            # Propagate never_match information to the whole group
+                            never_match_array[to_unite.detach().cpu().numpy()] = never_0 | never_1
+                            
+                        # If we are uniting more than two strings, we can eliminate
+                        # some redundant matches in the queue
+                        if not singletons:
+                            # Removed queued matches that are now in the same group
+                            matches = matches[group_ids[matches[:,0]] != group_ids[matches[:,1]]]
+
+                        if return_united:
+                            match_record = np.empty(4,dtype=int)
+                            match_record[:2] = match_pair.cpu().numpy().ravel()
+                            match_record[2] = self.counts[g0].sum().item()
+                            match_record[3] = self.counts[g1].sum().item()
+                            
+                            united.append(match_record)
+                    else:
+                        # Remove queued matches connecting these groups
+                        matches = matches[torch.isin(group_ids[matches[:,0]],g,invert=True) \
+                                            | torch.isin(group_ids[matches[:,1]],g,invert=True)]
+
+                    # Update progress bar
+                    p_bar.update(n_matches - matches.shape[0])
+                    n_matches = matches.shape[0]
+
+        predicted_grouping = self.ids_to_group(group_ids)
+
+        if always_match is not None:
+            predicted_grouping = predicted_grouping.unite(always_grouping)
+
+        if return_united:
+            united_df = pd.DataFrame(np.vstack(united),columns=['i','j','n_i','n_j'])
+            united_df = pd.merge(united_df,cos_scores_df,how='inner',on=['i','j'])
+            united_df['score'] = self.score_model(
+                                    torch.tensor(united_df['cos'].values).to(self.device)
+                                    ).cpu().numpy()
+            
+            united_df = united_df.drop('cos',axis=1)
+            
+            for c in ['i','j']:
+                united_df[c] = [self.strings[i] for i in united_df[c]]
+
+            if always_match is not None:
+                united_df['always_match'] = [always_grouping[i] == always_grouping[j] 
+                                            for i,j in united_df[['i','j']].values]
+
+            return predicted_grouping,united_df
+            
+        else:
+
+            return predicted_grouping
+
+    @torch.no_grad()
+    def unite_nearest(self,target_strings,threshold=0,always_grouping=None,progress_bar=True,batch_size=64):
+        """
+        Unite embedding strings with each string's most similar target string.
+
+        - "always_grouping" will be used to inialize the group_ids before uniting new matches
+        - "theshold" sets the minimimum match similarity required between a string and target string
+          for the string to be matched. (i.e., setting theshold=0 will result in every embedding
+          string to be matched its nearest target string, while setting threshold=0.9 will leave
+          strings that have similarity<0.9 with their nearest target string unaffected)
+
+        returns: MatchGroups object
+        """
+
+        if always_grouping is not None:
+            # self = self.embed(always_grouping)
+            group_ids = self._group_to_ids(always_grouping)
+        else:
+            group_ids = torch.arange(len(self)).to(self.device)
+
+        V = self.V
+        cos_threshold = self.score_model.score_to_cos(threshold)
+
+        seed_ids = torch.tensor([self.string_map[s] for s in target_strings]).to(self.device)
+        V_seed = V[seed_ids]
+        g_seed = group_ids[seed_ids]
+        is_seed = torch.zeros(V.shape[0],dtype=torch.bool).to(self.device)
+        is_seed[g_seed] = True
+
+        for batch_start in tqdm(range(0,len(self),batch_size),
+                                    delay=1,desc='Predicting matches',disable=not progress_bar):
+
+            batch_slice = slice(batch_start,batch_start+batch_size)
+
+            batch_cos = V[batch_slice]@V_seed.T
+
+            max_cos,max_seed = torch.max(batch_cos,dim=1)
+
+            # Get batch index locations where score > threshold
+            batch_i = torch.nonzero(max_cos > cos_threshold)
+
+            if len(batch_i):
+                # Drop target strings from matches (otherwise numerical precision
+                # issues can allow target strings to match to other strings)
+                batch_i = batch_i[~is_seed[batch_slice][batch_i]]
+
+                if len(batch_i):
+                    # Get indices of matched strings
+                    i = batch_i + batch_start
+
+                    # Assign matched strings to the target string's group
+                    group_ids[i] = g_seed[max_seed[batch_i]]
+
+        return self._ids_to_group(group_ids)
+
+    @torch.no_grad()
+    def score_pairs(self,string_pairs,batch_size=64,progress_bar=True):
+        string_pairs = np.array(string_pairs)
+
+        scores = []
+        for batch_start in tqdm(range(0,string_pairs.shape[0],batch_size),desc='Scoring pairs',disable=not progress_bar):
+
+            V0 = self[string_pairs[batch_start:batch_start+batch_size,0]].V
+            V1 = self[string_pairs[batch_start:batch_start+batch_size,1]].V
+
+            batch_cos = (V0*V1).sum(dim=1).ravel()
+            batch_scores = self.score_model(batch_cos)
+            
+            scores.append(batch_scores.cpu().numpy())
+
+        return np.concatenate(scores)
+
+    @torch.no_grad()
+    def _batch_scores(self,group_ids,batch_start,batch_size,
+                            is_match=None,
+                            min_score=None,max_score=None,
+                            min_loss=None,max_loss=None):
+
+        strings = self.strings
+        V = self.V
+        w = self.w
+
+        # Create simple slice objects to avoid creating copies with advanced indexing
+        i_slice = slice(batch_start,batch_start+batch_size)
+        j_slice = slice(batch_start+1,None)
+
+        X = V[i_slice]@V[j_slice].T
+        Y = (group_ids[i_slice,None] == group_ids[None,j_slice]).float()
+        if w is not None:
+            W = w[i_slice,None]*w[None,j_slice]
+        else:
+            W = None
+
+        scores = self.score_model(X)
+        loss = self.score_model.loss(X,Y,weights=W)
+
+        # Search upper diagonal entries only
+        # (note j_slice starting index is offset by one)
+        scores = torch.triu(scores)
+
+        # Filter by match type
+        if is_match is not None:
+            if is_match:
+                scores *= Y
+            else:
+                scores *= (1 - Y)
+
+        # Filter by min score
+        if min_score is not None:
+            scores *= (scores >= min_score)
+
+        # Filter by max score
+        if max_score is not None:
+            scores *= (scores <= max_score)
+
+        # Filter by min loss
+        if min_loss is not None:
+            scores *= (loss >= min_loss)
+
+        # Filter by max loss
+        if max_loss is not None:
+            scores *= (loss <= max_loss)
+
+        # Collect scored pairs
+        i,j = torch.nonzero(scores,as_tuple=True)
+
+        pairs = np.hstack([
+                                strings[i.cpu().numpy() + batch_start][:,None],
+                                strings[j.cpu().numpy() + (batch_start + 1)][:,None]
+                                ])
+
+        pair_groups = np.hstack([
+                                strings[group_ids[i + batch_start].cpu().numpy()][:,None],
+                                strings[group_ids[j + (batch_start + 1)].cpu().numpy()][:,None]
+                                ])
+
+        pair_scores = scores[i,j].cpu().numpy()
+        pair_losses = loss[i,j].cpu().numpy()
+
+        return pairs,pair_groups,pair_scores,pair_losses
+
+    def iter_scores(self,grouping=None,batch_size=64,progress_bar=True,**kwargs):
+
+        if grouping is not None:
+            self = self.embed(grouping)
+            group_ids = self._group_to_ids(grouping)
+        else:
+            group_ids = torch.arange(len(self)).to(self.device)
+
+        for batch_start in tqdm(range(0,len(self),batch_size),desc='Scoring pairs',disable=not progress_bar):
+            pairs,pair_groups,scores,losses = self._batch_scored_pairs(self,group_ids,batch_start,batch_size,**kwargs)
+            for (s0,s1),(g0,g1),score,loss in zip(pairs,pair_groups,scores,losses):
+                yield {
+                        'string0':s0,
+                        'string1':s1,
+                        'group0':g0,
+                        'group1':g1,
+                        'score':score,
+                        'loss':loss,
+                        }
+
+
+def load_embeddings(f):
+    """
+    Load embeddings from custom zipped archive format
+    """
+    with ZipFile(f,'r') as zip:
+        score_model = pickle.loads(zip.read('score_model.pkl'))
+        weighting_function = pickle.loads(zip.read('weighting_function.pkl'))
+        strings_df = pd.read_csv(zip.open('strings.csv'),na_filter=False)
+        V = np.load(zip.open('V.npy'))
+
+        return Embeddings(
+                            strings=strings_df['string'].values,
+                            counts=torch.tensor(strings_df['count'].values),
+                            score_model=score_model,
+                            weighting_function=weighting_function,
+                            V=torch.tensor(V)
+                            )
+

match_groups.py

@@ -0,0 +1,870 @@
+import os
+from pathlib import Path
+from collections import Counter, defaultdict
+from itertools import islice
+import pandas as pd
+import numpy as np
+import networkx as nx
+import matplotlib.pyplot as plt
+import matplotlib as mplt
+
+MAX_STR = 50
+
+
+class MatchGroups():
+    """A class for grouping strings based on set membership.  Supports splitting and uniting of groups."""
+
+    def __init__(self, strings=None):
+        """
+        Initialize MatchGroups object.
+
+        Parameters
+        ----------
+        strings : list, optional
+            List of strings to add to the match groups object, by default None
+        """
+        self.counts = Counter()
+        self.labels = {}
+        self.groups = {}
+
+        if strings is not None:
+            self.add_strings(strings, inplace=True)
+
+    def __len__(self):
+        """Return the number of strings in the match groups object."""
+        return len(self.labels)
+
+    def __repr__(self):
+        """Return a string representation of the MatchGroups object."""
+        return f'<nama.MatchGroups containing {len(self)} strings in {len(self.groups)} groups>'
+
+    def __str__(self):
+        """Return a string representation of the groups of a MatchGroups object."""
+        output = self.__repr__()
+        remaining = MAX_STR
+        for group in self.groups.values():
+            for s in group:
+                if remaining:
+                    output += '\n' + s
+                    remaining -= 1
+                else:
+                    output += f'...\n(Output truncated at {MAX_STR} strings)'
+                    return output
+
+            output += '\n'
+
+        return output
+
+    def __contains__(self, s):
+        """Return True if string is in the match groups object, False otherwise."""
+        return s in self.labels
+
+    def __getitem__(self, strings):
+        """Return the group label for a single string or a list of strings."""
+        if isinstance(strings, str):
+            return self.labels[strings]
+        else:
+            return [self.labels[s] for s in strings]
+
+    def __add__(self, match_obj):
+        """Add two match groups objects together and return the result."""
+        result = self.add_strings(match_obj)
+        result.unite(match_obj, inplace=True)
+
+        return result
+
+    def items(self):
+        """Return an iterator of strings and their group labels."""
+        for i, g in self.labels.items():
+            yield i, g
+
+    def copy(self):
+        """Return a copy of the MatchGroups object."""
+        new_match_obj = MatchGroups()
+        new_match_obj.counts = self.counts.copy()
+        new_match_obj.labels = self.labels.copy()
+        new_match_obj.groups = self.groups.copy()
+
+        return new_match_obj
+
+    def strings(self):
+        """Return a list of strings in the match groups object. Order is not guaranteed."""
+        return list(self.labels.keys())
+
+    def matches(self, string):
+        """Return the group of strings that match the given string."""
+        return self.groups[self.labels[string]]
+
+    def add_strings(self, arg, inplace=False):
+        """Add new strings to the match groups object.
+
+        Parameters
+        ----------
+        arg : str, Counter, MatchGroups, Iterable
+            String or group of strings to add to the match groups object
+        inplace : bool, optional
+            If True, add strings to the existing MatchGroups object, by default False
+
+        Returns
+        -------
+        MatchGroups
+            The updated MatchGroups object
+        """
+        if isinstance(arg, str):
+            counts = {arg: 1}
+
+        elif isinstance(arg, Counter):
+            counts = arg
+
+        elif isinstance(arg, MatchGroups):
+            counts = arg.counts
+
+        elif hasattr(arg, '__next__') or hasattr(arg, '__iter__'):
+            counts = Counter(arg)
+
+        if not inplace:
+            self = self.copy()
+
+        for s in counts.keys():
+            if s not in self.labels:
+                self.labels[s] = s
+                self.groups[s] = [s]
+
+        self.counts += counts
+
+        return self
+
+    def drop(self, strings, inplace=False):
+        """Remove strings from the match groups object.
+
+        Parameters
+        ----------
+        strings : list or str
+            String or list of strings to remove from the match groups object
+        inplace : bool, optional
+            If True, remove strings from the existing MatchGroups object, by default False
+
+        Returns
+        -------
+        MatchGroups
+            The updated MatchGroups object
+        """
+        if isinstance(strings, str):
+            strings = [strings]
+
+        strings = set(strings)
+
+        if not inplace:
+            self = self.copy()
+
+        # Remove strings from their groups
+        affected_group_labels = {self[s] for s in strings}
+        for old_label in affected_group_labels:
+            old_group = self.groups[old_label]
+            new_group = [s for s in old_group if s not in strings]
+
+            if new_group:
+                counts = self.counts
+                new_label = min((-counts[s], s) for s in new_group)[1]
+
+                if new_label != old_label:
+                    del self.groups[old_label]
+
+                self.groups[new_label] = new_group
+
+                for s in new_group:
+                    self.labels[s] = new_label
+            else:
+                del self.groups[old_label]
+
+        # Remove strings from counts and labels
+        for s in strings:
+            del self.counts[s]
+            del self.labels[s]
+
+        return self
+
+    def keep(self, strings, inplace=False):
+        """Drop all strings from the match groups object except the passed strings.
+
+        Parameters
+        ----------
+        strings : list
+            List of strings to keep in the match groups object
+        inplace : bool, optional
+            If True, drop strings from the existing MatchGroups object, by default False
+
+        Returns
+        -------
+        MatchGroups
+            The updated MatchGroups object
+        """
+        strings = set(strings)
+
+        to_drop = [s for s in self.strings() if s not in strings]
+
+        return self.drop(to_drop, inplace=inplace)
+
+    def _unite_strings(self, strings):
+        """
+        Unite strings in the match groups object without checking argument type.
+        Intended as a low-level function called by self.unite()
+
+        Parameters
+        ----------
+        strings : list
+            List of strings to unite in the match groups object
+
+        Returns
+        -------
+        None
+        """
+        strings = {s for s in strings if s in self.labels}
+
+        if len(strings) > 1:
+
+            # Identify groups that will be united
+            old_labels = set(self[strings])
+
+            # Only need to do the merge if the strings span multiple groups
+            if len(old_labels) > 1:
+
+                # Identify the new group label
+                counts = self.counts
+                new_label = min((-counts[s], s) for s in old_labels)[1]
+
+                # Identify the groups which need to be modified
+                old_labels.remove(new_label)
+
+                for old_label in old_labels:
+                    # Update the string group labels
+                    for s in self.groups[old_label]:
+                        self.labels[s] = new_label
+
+                    # Update group dict
+                    self.groups[new_label] = self.groups[new_label] + \
+                        self.groups[old_label]
+                    del self.groups[old_label]
+
+    def unite(self, arg, inplace=False, **kwargs):
+        """
+        Merge groups containing the passed strings. Groups can be passed as:
+        - A list of strings to unite
+        - A nested list to unite each set of strings
+        - A dictionary mapping strings to labels to unite by label
+        - A function mapping strings to labels to unite by label
+        - A MatchGroups instance to unite by MatchGroups groups
+
+        Parameters
+        ----------
+        arg : list, dict, function or MatchGroups instance
+            Argument representing the strings or labels to merge.
+        inplace : bool, optional
+            Whether to perform the operation in place or return a new MatchGroups.
+        kwargs : dict, optional
+            Additional arguments to be passed to predict_groupings method if arg
+            is a similarity model with a predict_groupings method.
+
+        Returns
+        -------
+        MatchGroups
+            The updated MatchGroups object. If `inplace` is True, the updated object
+            is returned, else a new MatchGroups object with the updates is returned.
+        """
+
+        if not inplace:
+            self = self.copy()
+
+        if isinstance(arg, str):
+            raise ValueError('Cannot unite a single string')
+
+        elif isinstance(arg, MatchGroups):
+            self.unite(arg.groups.values(), inplace=True)
+
+        elif hasattr(arg, 'unite_similar'):
+            # Unite can accept a similarity model if it has a unite_similar
+            # method
+            self.unite(arg.unite_similar(self, **kwargs))
+
+        elif callable(arg):
+            # Assume arg is a mapping from strings to labels and unite by label
+            groups = {s: arg(s) for s in self.strings()}
+            self.unite(groups, inplace=True)
+
+        elif isinstance(arg, dict):
+            # Assume arg is a mapping from strings to labels and unite by label
+            # groups = {label:[] for label in arg.values()}
+            groups = defaultdict(list)
+            for string, label in arg.items():
+                groups[label].append(string)
+
+            for group in groups.values():
+                self._unite_strings(group)
+
+        elif hasattr(arg, '__next__'):
+            # Assume arg is an iterator of groups to unite
+            # (This needs to be checked early to avoid consuming the first group)
+            for group in arg:
+                self._unite_strings(group)
+
+        elif all(isinstance(s, str) for s in arg):
+            # Main case: Unite group of strings
+            self._unite_strings(arg)
+
+        elif hasattr(arg, '__iter__'):
+            # Assume arg is an iterable of groups to unite
+            for group in arg:
+                self._unite_strings(group)
+
+        else:
+            raise ValueError('Unknown input type')
+
+        if not inplace:
+            return self
+
+    def split(self, strings, inplace=False):
+        """
+        Split strings into singleton groups. Strings can be passed as:
+        - A single string to isolate into a singleton group
+        - A list or iterator of strings to split
+
+        Parameters
+        ----------
+        strings : str or list of str
+            The string(s) to split into singleton groups.
+        inplace : bool, optional
+            Whether to perform the operation in place or return a new MatchGroups.
+
+        Returns
+        -------
+        MatchGroups
+            The updated MatchGroups object. If `inplace` is True, the updated object
+            is returned, else a new MatchGroups object with the updates is returned.
+        """
+        if not inplace:
+            self = self.copy()
+
+        if isinstance(strings, str):
+            strings = [strings]
+
+        strings = set(strings)
+
+        # Remove strings from their groups
+        affected_group_labels = {self[s] for s in strings}
+        for old_label in affected_group_labels:
+            old_group = self.groups[old_label]
+            if len(old_group) > 1:
+                new_group = [s for s in old_group if s not in strings]
+                if new_group:
+                    counts = self.counts
+                    new_label = min((-counts[s], s) for s in new_group)[1]
+
+                    if new_label != old_label:
+                        del self.groups[old_label]
+
+                    self.groups[new_label] = new_group
+
+                    for s in new_group:
+                        self.labels[s] = new_label
+
+        # Update labels and add singleton groups
+        for s in strings:
+            self.labels[s] = s
+            self.groups[s] = [s]
+
+        return self
+
+    def split_all(self, inplace=False):
+        """
+        Split all strings into singleton groups.
+
+        Parameters
+        ----------
+        inplace : bool, optional
+            Whether to perform the operation in place or return a new MatchGroups.
+
+        Returns
+        -------
+        MatchGroups
+            The updated MatchGroups object. If `inplace` is True, the updated object
+            is returned, else a new MatchGroups object with the updates is returned.
+        """
+        if not inplace:
+            self = self.copy()
+
+        self.labels = {s: s for s in self.strings()}
+        self.groups = {s: [s] for s in self.strings()}
+
+        return self
+
+    def separate(
+    self,
+    strings,
+    similarity_model,
+    inplace=False,
+    threshold=0,
+     **kwargs):
+        """
+        Separate the strings in according to the prediction of the similarity_model.
+
+        Parameters
+        ----------
+        strings: list
+            List of strings to be separated.
+        similarity_model: Model
+            Model used to predict similarity between strings.
+        inplace: bool, optional
+            If True, the separation operation is performed in-place. Otherwise, a copy is created.
+        threshold: float, optional
+            Threshold value for prediction.
+        kwargs: dict, optional
+            Additional keyword arguments passed to the prediction function.
+
+        Returns
+        -------
+        self: MatchGroups
+            Returns the MatchGroups object after the separation operation.
+
+        """
+        if not inplace:
+            self = self.copy()
+
+        # Identify which groups contain the strings to separate
+        group_map = defaultdict(list)
+        for s in set(strings):
+            group_map[self[s]].append(s)
+
+        for g, g_sep in group_map.items():
+
+            # If group contains strings to separate...
+            if len(g_sep) > 1:
+                group_strings = self.groups[g]
+
+                # Split the group strings
+                self.split(group_strings, inplace=True)
+
+                # Re-unite with new prediction that enforces separation
+                try:
+                    embeddings = similarity_model[group_strings]
+                except Exception as e:
+                    print(f'{g=} {g_sep} {group_strings}')
+                    raise e
+                predicted = embeddings.predict(
+                                threshold=threshold,
+                                separate_strings=strings,
+                                **kwargs)
+
+                self.unite(predicted, inplace=True)
+
+        return self
+
+    # def refine(self,similarity_model)
+
+    def top_scored_pairs_df(self, similarity_model,
+                                n=10000, buffer_n=100000,
+                                by_group=True,
+                                sort_by=['impact', 'score'], ascending=False,
+                                skip_pairs=None, **kwargs):
+        """
+        Return the DataFrame containing the n most important pairs of strings, according to the score generated by the `similarity_model`.
+
+        Parameters
+        ----------
+        similarity_model: Model
+            Model used to predict similarity between strings.
+        n: int, optional
+            Number of most important pairs to return. Default is 10000.
+        buffer_n: int, optional
+            Size of buffer to iterate through the scored pairs. Default is 100000.
+        by_group: bool, optional
+            If True, only the most important pair will be returned for each unique group combination.
+        sort_by: list, optional
+            A list of column names by which to sort the dataframe. Default is ['impact','score'].
+        ascending: bool, optional
+            Whether the sort order should be ascending or descending. Default is False.
+        skip_pairs: list, optional
+            List of string pairs to ignore when constructing the ranking.
+            If by_group=True, any group combination represented in the skip_pairs list will be ignored
+        kwargs: dict, optional
+            Additional keyword arguments passed to the `iter_scored_pairs` function.
+
+        Returns
+        -------
+        top_df: pandas.DataFrame
+            The DataFrame containing the n most important pairs of strings.
+
+        """
+
+        top_df = pd.DataFrame(
+    columns=[
+        'string0',
+        'string1',
+        'group0',
+        'group1',
+        'impact',
+        'score',
+         'loss'])
+        pair_iterator = similarity_model.iter_scored_pairs(self, **kwargs)
+
+        def group_size(g):
+            return len(self.groups[g])
+
+        if skip_pairs is not None:
+            if by_group:
+                skip_pairs = {tuple(sorted([self[s0], self[s1]]))
+                                    for s0, s1 in skip_pairs}
+            else:
+                skip_pairs = {tuple(sorted([s0, s1])) for s0, s1 in skip_pairs}
+
+        while True:
+            df = pd.DataFrame(islice(pair_iterator, buffer_n))
+
+            if len(df):
+                for i in 0, 1:
+                    df[f'group{i}'] = [self[s] for s in df[f'string{i}']]
+                df['impact'] = df['group0'].apply(
+                    group_size) * df['group1'].apply(group_size)
+
+                if by_group:
+                    df['group_pair'] = [tuple(sorted([g0, g1])) for g0, g1 in df[[
+                                              'group0', 'group1']].values]
+
+                if skip_pairs:
+                    if by_group:
+                        df = df[~df['group_pair'].isin(skip_pairs)]
+                    else:
+                        string_pairs = [tuple(sorted([s0, s1])) for s0, s1 in df[[
+                                              'string0', 'string1']].values]
+                        df = df[~string_pairs.isin(skip_pairs)]
+
+                if len(df):
+                    top_df = pd.concat([top_df, df]) \
+                                        .sort_values(sort_by, ascending=ascending)
+
+                    if by_group:
+                        top_df = top_df \
+                                        .groupby('group_pair') \
+                                        .first() \
+                                        .reset_index()
+
+                    top_df = top_df \
+                                        .sort_values(sort_by, ascending=ascending) \
+                                        .head(n)
+            else:
+                break
+
+        if len(top_df) and by_group:
+            top_df = top_df \
+                        .drop('group_pair', axis=1) \
+                        .reset_index()
+
+        return top_df
+
+    def reset_counts(self, inplace=False):
+        """
+        Reset the counts of strings in the MatchGroups object.
+
+        Parameters
+        ----------
+        inplace: bool, optional
+            If True, the operation is performed in-place. Otherwise, a copy is created.
+
+        Returns
+        -------
+        self: MatchGroups
+            Returns the MatchGroups object after the reset operation.
+
+        """
+        if not inplace:
+            self = self.copy()
+
+        self.counts = Counter(self.strings())
+
+        return self
+
+    def to_df(self, singletons=True, sort_groups=True):
+        """
+        Convert the match groups object to a dataframe with string, count and group columns.
+
+        Parameters
+        ----------
+        singletons: bool, optional
+            If True, the resulting DataFrame will include singleton groups. Default is True.
+        ...
+
+        Returns
+        -------
+        df: pandas.DataFrame
+            The resulting DataFrame.
+            """
+        strings = self.strings()
+
+        if singletons:
+            df = pd.DataFrame([(s, self.counts[s], self.labels[s]) for s in strings],
+                                columns=['string', 'count', 'group'])
+        else:
+            df = pd.DataFrame([(s, self.counts[s], self.labels[s]) for s in strings
+                                if len(self.groups[self[s]]) > 1],
+                                columns=['string', 'count', 'group'])
+        if sort_groups:
+            df['group_count'] = df.groupby('group')['count'].transform('sum')
+            df = df.sort_values(['group_count', 'group', 'count', 'string'], ascending=[
+                                False, True, False, True])
+            df = df.drop('group_count', axis=1)
+            df = df.reset_index(drop=True)
+
+        return df
+
+    def to_csv(self, filename, singletons=True, **pandas_args):
+        """
+        Save the match groups object as a csv file with string, count and group columns.
+
+        Parameters
+        ----------
+        filename : str
+            Path to file to save the data.
+        singletons : bool, optional
+            If True, include singleton groups in the saved file, by default True.
+        pandas_args : dict
+            Additional keyword arguments to pass to the pandas.DataFrame.to_csv method.
+        """
+        df = self.to_df(singletons=singletons)
+        df.to_csv(filename, index=False, **pandas_args)
+
+    def merge_dfs(self, left_df, right_df, how='inner',
+                    on=None, left_on=None, right_on=None,
+                    group_column_name='match_group', suffixes=('_x', '_y'),
+                    **merge_args):
+        """
+        Replicated pandas.merge() functionality, except that dataframes are merged by match group instead of directly on the strings in the "on" columns.
+
+        Parameters
+        ----------
+        left_df : pandas.DataFrame
+            The left dataframe to merge.
+        right_df : pandas.DataFrame
+            The right dataframe to merge.
+        how : str, optional
+            How to merge the dataframes. Possible values are 'left', 'right', 'outer', 'inner', by default 'inner'.
+        on : str, optional
+            Columns in both left and right dataframes to merge on.
+        left_on : str, optional
+            Columns in the left dataframe to merge on.
+        right_on : str, optional
+            Columns in the right dataframe to merge on.
+        group_column_name : str, optional
+            Column name for the merged match group, by default 'match_group'.
+        suffixes : tuple of str, optional
+            Suffix to apply to overlapping column names in the left and right dataframes, by default ('_x','_y').
+        **merge_args : dict
+            Additional keyword arguments to pass to the pandas.DataFrame.merge method.
+
+        Returns
+        -------
+        pandas.DataFrame
+            The merged dataframe.
+
+        Raises
+        ------
+        ValueError
+            If 'on', 'left_on', and 'right_on' are all None.
+        ValueError
+            If `group_column_name` already exists in one of the dataframes.
+        """
+
+        if ((left_on is None) or (right_on is None)) and (on is None):
+            raise ValueError('Must provide column(s) to merge on')
+
+        left_df = left_df.copy()
+        right_df = right_df.copy()
+
+        if on is not None:
+            left_on = on + suffixes[0]
+            right_on = on + suffixes[1]
+
+            left_df = left_df.rename(columns={on:left_on})
+            right_df = right_df.rename(columns={on:right_on})
+
+        group_map = lambda s: self[s] if s in self.labels else np.nan
+
+        left_group = left_df[left_on].apply(group_map)
+        right_group = right_df[right_on].apply(group_map)
+
+        if group_column_name:
+            if group_column_name in list(left_df.columns) + list(right_df.columns):
+                raise ValueError('f{group_column_name=} already exists in one of the dataframes.')
+            else:
+                left_df[group_column_name] = left_group
+
+        merged_df = pd.merge(left_df,right_df,left_on=left_group,right_on=right_group,how=how,suffixes=suffixes,**merge_args)
+
+        merged_df = merged_df[[c for c in merged_df.columns if c in list(left_df.columns) + list(right_df.columns)]]
+
+        return merged_df
+
+
+def from_df(
+        df,
+        match_format='detect',
+        pair_columns=[
+            'string0',
+            'string1'],
+    string_column='string',
+    group_column='group',
+        count_column='count'):
+    """
+    Construct a new match groups object from a pandas DataFrame.
+
+    Parameters
+    ----------
+    df : pandas.DataFrame
+        The input dataframe.
+    match_format : str, optional
+        The format of the dataframe, by default "detect".
+        It can be one of ['unmatched', 'groups', 'pairs', 'detect'].
+    pair_columns : list of str, optional
+        The columns names containing the string pairs, by default ['string0','string1'].
+    string_column : str, optional
+        The column name containing the strings, by default 'string'.
+    group_column : str, optional
+        The column name containing the groups, by default 'group'.
+    count_column : str, optional
+        The column name containing the counts, by default 'count'.
+
+    Returns
+    -------
+    MatchGroups
+        The constructed MatchGroups object.
+
+    Raises
+    ------
+    ValueError
+        If the input `match_format` is not one of ['unmatched', 'groups', 'pairs', 'detect'].
+    ValueError
+        If the `match_format` is 'detect' and the input dataframe format could not be inferred.
+
+    Notes
+    -----
+    The function accepts two formats of the input dataframe:
+
+        - "groups": The standard format for a match groups object dataframe. It includes a
+          string column, and a "group" column that contains group labels, and an
+          optional "count" column. These three columns completely describe a
+          match groups object, allowing lossless match groups object -> dataframe -> match groups object
+          conversion (though the specific group labels in the dataframe will be
+          ignored and rebuilt in the new match groups object).
+
+        - "pairs": The dataframe includes two string columns, and each row indicates
+          a link between a pair of strings. A new match groups object will be constructed by
+          uniting each pair of strings.
+    """
+
+    if match_format not in ['unmatched', 'groups', 'pairs', 'detect']:
+        raise ValueError(
+            'match_format must be one of "unmatched", "groups", "pairs", or "detect"')
+
+    # Create an empty match groups object
+    match_obj = MatchGroups()
+
+    if match_format == 'detect':
+        if (string_column in df.columns):
+            if group_column is None:
+                match_format = 'unmatched'
+            elif (group_column in df.columns):
+                match_format = 'groups'
+        elif set(df.columns) == set(pair_columns):
+            match_format = 'pairs'
+
+    if match_format == 'detect':
+        raise ValueError('Could not infer valid dataframe format from input')
+
+    if count_column in df.columns:
+        counts = df[count_column].values
+    else:
+        counts = np.ones(len(df))
+
+    if match_format == 'unmatched':
+        strings = df[string_column].values
+
+        # Build the match groups object
+        match_obj.counts = Counter({s: int(c) for s, c in zip(strings, counts)})
+        match_obj.labels = {s: s for s in strings}
+        match_obj.groups = {s: [s] for s in strings}
+
+    elif match_format == 'groups':
+
+        strings = df[string_column].values
+        group_ids = df[group_column].values
+
+        # Sort by group and string count
+        g_sort = np.lexsort((counts, group_ids))
+        group_ids = group_ids[g_sort]
+        strings = strings[g_sort]
+        counts = counts[g_sort]
+
+        # Identify group boundaries and split locations
+        split_locs = np.nonzero(group_ids[1:] != group_ids[:-1])[0] + 1
+
+        # Get grouped strings as separate arrays
+        groups = np.split(strings, split_locs)
+
+        # Build the match groups object
+        match_obj.counts = Counter({s: int(c) for s, c in zip(strings, counts)})
+        match_obj.labels = {s: g[-1] for g in groups for s in g}
+        match_obj.groups = {g[-1]: list(g) for g in groups}
+
+    elif match_format == 'pairs':
+        # TODO: Allow pairs data to use counts
+        for pair_column in pair_columns:
+            match_obj.add_strings(df[pair_column].values, inplace=True)
+
+        # There are several ways to unite pairs
+        # Guessing it is most efficient to "group by" one of the string columns
+        groups = {s: pair[1] for pair in df[pair_columns].values for s in pair}
+
+        match_obj.unite(groups, inplace=True)
+
+    return match_obj
+
+
+def read_csv(
+        filename,
+        match_format='detect',
+        pair_columns=[
+            'string0',
+            'string1'],
+    string_column='string',
+    group_column='group',
+    count_column='count',
+        **pandas_args):
+    """
+    Read a csv file and construct a new match groups object.
+
+    Parameters
+    ----------
+    filename : str
+        The path to the csv file.
+    match_format : str, optional (default='detect')
+        One of "unmatched", "groups", "pairs", or "detect".
+    pair_columns : list of str, optional (default=['string0', 'string1'])
+        Two string columns to use if match_format='pairs'.
+    string_column : str, optional (default='string')
+        Column name for string values in match_format='unmatched' or 'groups'.
+    group_column : str, optional (default='group')
+        Column name for group values in match_format='groups'.
+    count_column : str, optional (default='count')
+        Column name for count values in match_format='unmatched' or 'groups'.
+    **pandas_args : optional
+        Optional arguments to pass to `pandas.read_csv`.
+
+    Returns
+    -------
+    MatchGroups
+        A new match groups object built from the csv file.
+    """
+    df = pd.read_csv(filename, **pandas_args, na_filter=False)
+    df = df.astype(str)
+
+    return from_df(df, match_format=match_format, pair_columns=pair_columns,
+                   string_column=string_column, group_column=group_column,
+                   count_column=count_column)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e112a851e5079096d2f0bab96d21dfb0dea8cd92a2f23fe9218cc65dd9777fbe
+size 499051193

scoring.py

@@ -0,0 +1,196 @@
+import pandas as pd
+import random
+
+
+def confusion_df(predicted_groupings, gold_groupings, use_counts=True):
+    """
+    Computes the confusion matrix dataframe for a predicted match groups object relative to a gold match groups object.
+
+    Parameters
+    ----------
+    predicted_groupings : MatchGroups
+        The predicted match groups object.
+    gold_groupings : MatchGroups
+        The gold match groups object.
+    use_counts : bool, optional
+        Use the count of each string. If False, the count is set to 1.
+
+    Returns
+    -------
+    df : pandas.DataFrame
+        Confusion matrix dataframe with columns 'TP', 'FP', 'TN', and 'FN'.
+    """
+
+    df = pd.merge(
+        predicted_groupings.to_df(),
+        gold_groupings.to_df().drop(
+            'count',
+            axis=1),
+        on='string',
+        suffixes=[
+            '_pred',
+            '_gold'])
+
+    if not use_counts:
+        df['count'] = 1
+
+    df['TP'] = (df.groupby(['group_pred', 'group_gold'])[
+                'count'].transform('sum') - df['count']) * df['count']
+    df['FP'] = (df.groupby('group_pred')['count'].transform(
+        'sum') - df['count']) * df['count'] - df['TP']
+    df['FN'] = (df.groupby('group_gold')['count'].transform(
+        'sum') - df['count']) * df['count'] - df['TP']
+    df['TN'] = (df['count'].sum() - df['count']) * \
+        df['count'] - df['TP'] - df['FP'] - df['FN']
+
+    return df
+
+
+def confusion_matrix(predicted_groupings, gold_groupings, use_counts=True):
+    """
+    Computes the confusion matrix for a predicted match groups object relative to a gold match groups object.
+
+    Parameters
+    ----------
+    predicted_groupings : MatchGroups
+        The predicted match groups object.
+    gold_groupings : MatchGroups
+        The gold match groups object.
+    use_counts : bool, optional
+        Use the count of each string. If False, the count is set to 1.
+
+    Returns
+    -------
+    confusion_matrix : dict
+        Dictionary with keys 'TP', 'FP', 'TN', and 'FN', representing the values in the confusion matrix.
+    """
+
+    df = confusion_df(predicted_groupings, gold_groupings, use_counts=use_counts)
+
+    return {c: df[c].sum() // 2 for c in ['TP', 'FP', 'TN', 'FN']}
+
+
+def score_predicted(
+        predicted_groupings,
+        gold_groupings,
+        use_counts=True,
+        drop_self_matches=True):
+    """
+    Computes the F1 score of a predicted match groups object relative to a gold match groups object
+    which is assumed to be correct.
+
+    Parameters
+    ----------
+    predicted_groupings : MatchGroups
+        The predicted match groups object .
+    gold_groupings : MatchGroups
+        The gold match groups object.
+    use_counts : bool, optional
+        Use the count of each string. If False, the count is set to 1.
+    drop_self_matches : bool, optional
+        Remove the matches between a string and itself.
+
+    Returns
+    -------
+    scores : dict
+        Dictionary with keys 'accuracy', 'precision', 'recall', 'F1', and 'coverage'.
+    """
+
+    scores = confusion_matrix(
+        predicted_groupings,
+        gold_groupings,
+        use_counts=use_counts)
+
+    n_scored = scores['TP'] + scores['TN'] + scores['FP'] + scores['FN']
+
+    if use_counts:
+        n_predicted = (sum(predicted_groupings.counts.values())**2 -
+                       sum(c**2 for c in predicted_groupings.counts.values())) / 2
+    else:
+        n_predicted = (len(predicted_groupings)**2
+                       - len(predicted_groupings)) / 2
+
+    scores['coverage'] = n_scored / n_predicted
+
+    if scores['TP']:
+        scores['accuracy'] = (scores['TP'] + scores['TN']) / n_scored
+        scores['precision'] = scores['TP'] / (scores['TP'] + scores['FP'])
+        scores['recall'] = scores['TP'] / (scores['TP'] + scores['FN'])
+        scores['F1'] = 2 * (scores['precision'] * scores['recall']) / \
+            (scores['precision'] + scores['recall'])
+
+    else:
+        scores['accuracy'] = 0
+        scores['precision'] = 0
+        scores['recall'] = 0
+        scores['F1'] = 0
+
+    return scores
+
+
+def split_on_groups(groupings, frac=0.5, seed=None):
+    """
+    Splits the match groups object into two parts by given fraction.
+
+    Parameters
+    ----------
+    groupings : MatchGroups
+        The match groups object to be split.
+    frac : float, optional
+        The fraction of groups to select.
+    seed : int, optional
+        Seed for the random number generator.
+
+    Returns
+    -------
+    groupings1, groupings2 : tuple of match groups objects
+        Tuple of two match groups objects.
+    """
+    if seed is not None:
+        random.seed(seed)
+
+    groups = list(groupings.groups.values())
+    random.shuffle(groups)
+
+    selected_groups = groups[:int(frac * len(groups))]
+    selected_strings = [s for group in selected_groups for s in group]
+
+    return groupings.keep(selected_strings), groupings.drop(selected_strings)
+
+
+def kfold_on_groups(groupings, k=4, shuffle=True, seed=None):
+    """
+    Perform K-fold cross validation on groups of strings.
+
+    Parameters
+    ----------
+    groupings : object
+        MatchGroups object to perform K-fold cross validation on.
+    k : int, optional
+        Number of folds to perform, by default 4.
+    shuffle : bool, optional
+        Whether to shuffle the groups before splitting, by default True.
+    seed : int, optional
+        Seed for the random number generator, by default None.
+
+    Yields
+    ------
+    tuple : MatchGroups, MatchGroups
+        A tuple of k match groups objects, the first for the training set and the second for the testing set for each fold.
+    """
+    if seed is not None:
+        random.seed(seed)
+
+    groups = list(groupings.groups.keys())
+
+    if shuffle:
+        random.shuffle(groups)
+    else:
+        groups = sorted(groups)
+
+    for fold in range(k):
+
+        fold_groups = groups[fold::k]
+        fold_strings = [s for g in fold_groups for s in groupings.groups[g]]
+
+        yield groupings.drop(fold_strings), groupings.keep(fold_strings)

scoring_model.py

@@ -0,0 +1,60 @@
+import torch
+import numpy as np
+
+class SimilarityScore(torch.nn.Module):
+    """
+    A trainable similarity scoring model that estimates the probability
+    of a match as the negative exponent of 1+cosine distance between
+    embeddings:
+        p(match|v_i,v_j) = exp(-alpha*(1-v_i@v_j))
+    """
+    def __init__(self,config,**kwargs):
+
+        super().__init__()
+
+        self.alpha = torch.nn.Parameter(torch.tensor(float(config.get("alpha"))))
+
+    def __repr__(self):
+        return f'<nama.ExpCosSimilarity with {self.alpha=}>'
+
+    def forward(self,X):
+        # Z is a scaled distance measure: Z=0 means that the score should be 1
+        Z = self.alpha*(1 - X)
+        return torch.clamp(torch.exp(-Z),min=0,max=1.0)
+
+    def loss(self,X,Y,weights=None,decay=1e-6,epsilon=1e-6):
+
+        Z = self.alpha*(1 - X)
+
+        # Put epsilon floor to prevent overflow/undefined results
+        # Z = torch.tensor([1e-2,1e-3,1e-6,1e-7,1e-8,1e-9])
+        # torch.log(1 - torch.exp(-Z))
+        # 1/(1 - torch.exp(-Z))
+        with torch.no_grad():
+            Z_eps_adjustment = torch.clamp(epsilon-Z,min=0)
+
+        Z += Z_eps_adjustment
+
+        # Cross entropy loss with a simplified and (hopefully) numerically appropriate formula
+        # TODO: Stick an epsilon in here to prevent nan?
+        loss = Y*Z - torch.xlogy(1-Y,-torch.expm1(-Z))
+        # loss = Y*Z - torch.xlogy(1-Y,1-torch.exp(-Z))
+
+        if weights is not None:
+            loss *= weights*loss
+
+        if decay:
+            loss += decay*self.alpha**2
+
+        return loss
+
+    def score_to_cos(self,score):
+        if score > 0:
+            return 1 + np.log(score)/self.alpha.item()
+        else:
+            return -99
+
+    def config_optimizer(self,lr=10):
+        optimizer = torch.optim.AdamW(self.parameters(),lr=lr,weight_decay=0)
+
+        return optimizer

similarity_model.py

@@ -0,0 +1,369 @@
+import random
+import pandas as pd
+import numpy as np
+from tqdm import tqdm
+from copy import copy,deepcopy
+from collections import Counter
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from transformers import get_cosine_schedule_with_warmup,get_linear_schedule_with_warmup, logging
+from transformers.modeling_utils import PreTrainedModel
+
+from .match_groups import MatchGroups
+from .scoring import score_predicted
+from .scoring_model import SimilarityScore
+from .embeddings import Embeddings
+from .embedding_model import EmbeddingModel
+from .configuration import SimilarityModelConfig
+logging.set_verbosity_error()
+
+
+class ExponentWeights():
+    def __init__(self, config,**kwargs):
+        self.exponent = config.get("weighting_exponent", 0.5)
+
+    def __call__(self,counts):
+        return counts**self.exponent
+
+
+class SimilarityModel(PreTrainedModel):
+    config_class = SimilarityModelConfig
+    """
+    A combined embedding/scorer model that produces Embeddings objects
+    as its primary output.
+
+    - train() jointly optimizes the embedding_model and score_model using
+      contrastive learning to learn from a training MatchGroups.
+    """
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        
+        self.embedding_model = EmbeddingModel(config.embedding_model_config, **kwargs)
+        self.score_model = SimilarityScore(config.score_model_config, **kwargs)
+        self.weighting_function = ExponentWeights(config.weighting_function_config, **kwargs)
+        
+        self.config = config
+        self.to(config.device)
+
+    def to(self,device):
+        super().to(device)
+        self.embedding_model.to(device)
+        self.score_model.to(device)
+        #self.device = device
+
+    def save(self,savefile):
+        torch.save({'metadata': self.config, 'state_dict': self.state_dict()}, savefile)
+
+    @torch.no_grad()
+    def embed(self,input,to=None,batch_size=64,progress_bar=True,**kwargs):
+        """
+        Construct an Embeddings object from input strings or a MatchGroups
+        """
+
+        if to is None:
+            to = self.device
+
+        if isinstance(input, MatchGroups):
+            strings = input.strings()
+            counts = torch.tensor([input.counts[s] for s in strings],device=self.device).float().to(to)
+
+        else:
+            strings = list(input)
+            counts = torch.ones(len(strings),device=self.device).float().to(to)
+
+        input_loader = DataLoader(strings,batch_size=batch_size,num_workers=0)
+
+        self.embedding_model.eval()
+
+        V = None
+        batch_start = 0
+        with tqdm(total=len(strings),delay=1,desc='Embedding strings',disable=not progress_bar) as pbar:
+            for batch_strings in input_loader:
+
+                v = self.embedding_model(batch_strings).detach().to(to)
+
+                if V is None:
+                    # Use v to determine dim and dtype of pre-allocated embedding tensor
+                    # (Pre-allocating avoids duplicating tensors with a big .cat() operation)
+                    V = torch.empty(len(strings),v.shape[1],device=to,dtype=v.dtype)
+
+                V[batch_start:batch_start+len(batch_strings),:] = v
+
+                pbar.update(len(batch_strings))
+                batch_start += len(batch_strings)
+
+        score_model = copy(self.score_model)
+        score_model.load_state_dict(self.score_model.state_dict())
+        score_model.to(to)
+
+        weighting_function = deepcopy(self.weighting_function)
+
+        return Embeddings(strings=strings,
+                            V=V.detach(),
+                            counts=counts.detach(),
+                            score_model=score_model,
+                            weighting_function=weighting_function,
+                            device=to)
+
+    def train(self,training_groupings,max_epochs=1,batch_size=8,
+                score_decay=0,regularization=0,
+                transformer_lr=1e-5,projection_lr=1e-5,score_lr=10,warmup_frac=0.1,
+                max_grad_norm=1,dropout=False,
+                validation_groupings=None,target='F1',restore_best=True,val_seed=None,
+                validation_interval=1000,early_stopping=True,early_stopping_patience=3,
+                verbose=False,progress_bar=True,
+                **kwargs):
+
+        """
+        Train the embedding_model and score_model to predict match probabilities
+        using the training_groupings as a source of "correct" matches.
+        Training algorithm uses contrastive learning with hard-positive
+        and hard-negative mining to fine tune the embedding model to place
+        matched strings near to each other in embedding space, while
+        simulataneously calibrating the score_model to predict the match
+        probabilities as a function of cosine distance
+        """
+
+        if validation_groupings is None:
+            early_stopping = False
+            restore_best = False
+
+        num_training_steps = max_epochs*len(training_groupings)//batch_size
+        num_warmup_steps = int(warmup_frac*num_training_steps)
+
+        if transformer_lr or projection_lr:
+            embedding_optimizer = self.embedding_model.config_optimizer(transformer_lr,projection_lr)
+            embedding_scheduler = get_cosine_schedule_with_warmup(
+                                        embedding_optimizer,
+                                        num_warmup_steps=num_warmup_steps,
+                                        num_training_steps=num_training_steps)
+        if score_lr:
+            score_optimizer = self.score_model.config_optimizer(score_lr)
+            score_scheduler = get_linear_schedule_with_warmup(
+                                        score_optimizer,
+                                        num_warmup_steps=num_warmup_steps,
+                                        num_training_steps=num_training_steps)
+
+        step = 0
+        self.history = []
+        self.val_scores = []
+        for epoch in range(max_epochs):
+
+            global_embeddings = self.embed(training_groupings)
+
+            strings = global_embeddings.strings
+            V = global_embeddings.V
+            w = global_embeddings.w
+
+            groups = torch.tensor([global_embeddings.string_map[training_groupings[s]] for s in strings],device=self.device)
+
+            # Normalize weights to make learning rates more general
+            if w is not None:
+                w = w/w.mean()
+
+            shuffled_ids = list(range(len(strings)))
+            random.shuffle(shuffled_ids)
+
+            if dropout:
+                self.embedding_model.train()
+            else:
+                self.embedding_model.eval()
+
+            for batch_start in tqdm(range(0,len(strings),batch_size),desc=f'training epoch {epoch}',disable=not progress_bar):
+
+                h = {'epoch':epoch,'step':step}
+
+                batch_i = shuffled_ids[batch_start:batch_start+batch_size]
+
+                # Recycle ids from the beginning to pad the last batch if necessary
+                if len(batch_i) < batch_size:
+                    batch_i = batch_i + shuffled_ids[:(batch_size-len(batch_i))]
+
+                """
+                Find highest loss match for each batch string (global search)
+
+                Note: If we compute V_i with dropout enabled, it will add noise
+                to the embeddings and prevent the same pairs from being selected
+                every time.
+                """
+                V_i = self.embedding_model(strings[batch_i])
+
+                # Update global embedding cache
+                V[batch_i,:] = V_i.detach()
+
+                with torch.no_grad():
+
+                    global_X = V_i@V.T
+                    global_Y = (groups[batch_i][:,None] == groups[None,:]).float()
+
+                    if w is not None:
+                        global_W = torch.outer(w[batch_i],w)
+                    else:
+                        global_W = None
+
+                # Train scoring model only
+                if score_lr:
+                    # Make sure gradients are enabled for score model
+                    self.score_model.requires_grad_(True)
+
+                    global_loss = self.score_model.loss(global_X,global_Y,weights=global_W,decay=score_decay)
+
+                    score_optimizer.zero_grad()
+                    global_loss.nanmean().backward()
+                    torch.nn.utils.clip_grad_norm_(self.score_model.parameters(),max_norm=max_grad_norm)
+
+                    score_optimizer.step()
+                    score_scheduler.step()
+
+                    h['score_lr'] = score_optimizer.param_groups[0]['lr']
+                    h['global_mean_cos'] = global_X.mean().item()
+                    try:
+                        h['score_alpha'] = self.score_model.alpha.item()
+                    except:
+                        pass
+
+                else:
+                    with torch.no_grad():
+                        global_loss = self.score_model.loss(global_X,global_Y)
+
+                h['global_loss'] = global_loss.detach().nanmean().item()
+
+                # Train embedding model
+                if (transformer_lr or projection_lr) and step <= num_warmup_steps + num_training_steps:
+
+                    # Turn off score model updating - only want to train embedding here
+                    self.score_model.requires_grad_(False)
+
+                    # Select hard training examples
+                    with torch.no_grad():
+                        batch_j = global_loss.argmax(dim=1).flatten()
+
+                        if w is not None:
+                            batch_W = torch.outer(w[batch_i],w[batch_j])
+                        else:
+                            batch_W = None
+
+                    # Train the model on the selected high-loss pairs
+                    V_j = self.embedding_model(strings[batch_j.tolist()])
+
+                    # Update global embedding cache
+                    V[batch_j,:] = V_j.detach()
+
+                    batch_X = V_i@V_j.T
+                    batch_Y = (groups[batch_i][:,None] == groups[batch_j][None,:]).float()
+                    h['batch_obs'] = len(batch_i)*len(batch_j)
+
+                    batch_loss = self.score_model.loss(batch_X,batch_Y,weights=batch_W)
+
+                    if regularization:
+                        # Apply Global Orthogonal Regularization from https://arxiv.org/abs/1708.06320
+                        gor_Y = (groups[batch_i][:,None] != groups[batch_i][None,:]).float()
+                        gor_n = gor_Y.sum()
+                        if gor_n > 1:
+                            gor_X = (V_i@V_i.T)*gor_Y
+                            gor_m1 = 0.5*gor_X.sum()/gor_n
+                            gor_m2 = 0.5*(gor_X**2).sum()/gor_n
+                            batch_loss += regularization*(gor_m1 + torch.clamp(gor_m2 - 1/self.embedding_model.d,min=0))
+
+                    h['batch_nan'] = torch.isnan(batch_loss.detach()).sum().item()
+
+                    embedding_optimizer.zero_grad()
+                    batch_loss.nanmean().backward()
+
+                    torch.nn.utils.clip_grad_norm_(self.parameters(),max_norm=max_grad_norm)
+
+                    embedding_optimizer.step()
+                    embedding_scheduler.step()
+
+                    h['transformer_lr'] = embedding_optimizer.param_groups[1]['lr']
+                    h['projection_lr'] = embedding_optimizer.param_groups[-1]['lr']
+
+                    # Save stats
+                    h['batch_loss'] = batch_loss.detach().mean().item()
+                    h['batch_pos_target'] = batch_Y.detach().mean().item()
+
+                self.history.append(h)
+                step += 1
+
+                if (validation_groupings is not None) and not (step % validation_interval):
+
+                    validation = len(self.validation_scores)
+                    val_scores = self.test(validation_groupings)
+                    val_scores['step'] = step - 1
+                    val_scores['epoch'] = epoch
+                    val_scores['validation'] = validation
+
+                    self.validation_scores.append(val_scores)
+
+                    # Print validation stats
+                    if verbose:
+                        print(f'\nValidation results at step {step} (current epoch {epoch})')
+                        for k,v in val_scores.items():
+                            print(f'    {k}: {v:.4f}')
+
+                        print(list(self.score_model.named_parameters()))
+
+                    # Update best saved model
+                    if restore_best:
+                        if val_scores[target] >= max(h[target] for h in self.validation_scores):
+                            best_state = deepcopy({
+                                            'state_dict':self.state_dict(),
+                                            'val_scores':val_scores
+                                            })
+
+                    if early_stopping and (validation - best_state['val_scores']['validation'] > early_stopping_patience):
+                        print(f'Stopping training ({early_stopping_patience} validation checks since best validation score)')
+                        break
+
+        if restore_best:
+            print(f"Restoring to best state (step {best_state['val_scores']['step']}):")
+            for k,v in best_state['val_scores'].items():
+                print(f'    {k}: {v:.4f}')
+
+            self.to('cpu')
+            self.load_state_dict(best_state['state_dict'])
+            self.to(self.device)
+
+        return pd.DataFrame(self.history)
+
+    def unite_similar(self,input,**kwargs):
+        embeddings = self.embed(input,**kwargs)
+        return embeddings.unite_similar(**kwargs)
+
+    def test(self,gold_groupings, threshold=0.5, **kwargs):
+        embeddings = self.embed(gold_groupings, **kwargs)
+
+        if (isinstance(threshold, float)):
+            predicted = embeddings.unite_similar(threshold=threshold, **kwargs)
+            scores = score_predicted(predicted, gold_groupings, use_counts=True)
+
+            return scores
+        
+        results = []
+        for thres in threshold:
+            predicted = embeddings.unite_similar(threshold=thres, **kwargs)
+
+            scores = score_predicted(predicted, gold_groupings, use_counts=True)
+            scores["threshold"] = thres
+            results.append(scores)
+
+        
+        return results
+
+
+
+def load_similarity_model(f,map_location='cpu',*args,**kwargs):
+    checkpoint = torch.load(f, map_location=map_location, **kwargs)
+    metadata = checkpoint['metadata']
+    state_dict = checkpoint['state_dict']
+
+    model = SimilarityModel(config=metadata)
+    model.load_state_dict(state_dict)
+
+    return model
+    #return torch.load(f,map_location=map_location,**kwargs)
+
+    
+
+