tmnt 0.7.0b20230910__py3-none-any.whl → 0.7.0b20230912__py3-none-any.whl

tmnt/data_loading.py

@@ -92,7 +92,7 @@ def get_llm_paired_dataloader(data_a, data_b, bow_vectorizer, llm_name, label_ma
 
 class StratifiedPairedLLMLoader():
 
-    def __init__(self, data_a, data_b, bow_vectorizer, llm_name, label_map, batch_size, max_len_a, max_len_b, device='cpu'):
+    def __init__(self, data_a, data_b, bow_vectorizer, llm_name, label_map, batch_size, max_len_a, max_len_b, num_batches=0, device='cpu'):
         self.data_a = data_a
         self.data_b = data_b
         self.bow_vectorizer = bow_vectorizer
@@ -102,7 +102,7 @@ class StratifiedPairedLLMLoader():
         self.max_len_a = max_len_a
         self.max_len_b = max_len_b
         self.device = device
-        self.num_batches = max(len(data_a), len(data_b)) // batch_size
+        self.num_batches = num_batches or max(len(data_a), len(data_b)) // batch_size
         self.stratified_sampler = StratifiedDualBatchSampler(np.array([label_map[l] for (l,_) in data_a]), 
                                                              np.array([label_map[l] for (l,_) in data_b]),
                                                              batch_size,
@@ -419,22 +419,22 @@ class StratifiedDualBatchSampler:
         self.shuffle = shuffle
         self.batch_size = batch_size
         self.num_batches = num_batches
-        counts_a = Counter(y_a) 
-        counts_b = Counter(y_b)
+        self.counts_a = Counter(y_a) 
+        self.counts_b = Counter(y_b)
         self.class_weights_a = [0] * (max(np.max(y_a), np.max(y_b)) + 1)
         self.class_weights_b = [0] * (max(np.max(y_a), np.max(y_b)) + 1)
-        for k in counts_a:
-            self.class_weights_a[k] = counts_a[k] / len(y_a)
-        for k in counts_b:
-            self.class_weights_b[k] = counts_b[k] / len(y_b)
+        for k in self.counts_a:
+            self.class_weights_a[k] = self.counts_a[k] / len(y_a)
+        for k in self.counts_b:
+            self.class_weights_b[k] = self.counts_b[k] / len(y_b)
         self.class_indices_a = [0] * (max(np.max(y_a), np.max(y_b)) + 1)
         self.class_indices_b = [0] * (max(np.max(y_b), np.max(y_a)) + 1)
         for i in range(len(self.class_indices_a)):
             self.class_indices_a[i] = list(np.where(y_a == i)[0])
         for i in range(len(self.class_indices_b)):
             self.class_indices_b[i] = list(np.where(y_b == i)[0])
-        self.a_only = counts_a.keys() - counts_b.keys()
-        self.b_only = counts_b.keys() - counts_a.keys()
+        self.a_only = self.counts_a.keys() - self.counts_b.keys()
+        self.b_only = self.counts_b.keys() - self.counts_a.keys()
         self.use_with_replacement = (self.batch_size > len(self.class_weights_a))
         
     def _pop_leave_last(self, li):
@@ -451,14 +451,14 @@ class StratifiedDualBatchSampler:
         for i in range(self.num_batches):
             if i % 2 == 0:
                 classes_a = list(WeightedRandomSampler(self.class_weights_a, self.batch_size, replacement=self.use_with_replacement))
-                b_list = list(self.b_only)
+                b_list = list(self.counts_b)
                 random.shuffle(b_list)
                 classes_b = [ self._pop_leave_last(b_list) if a in self.a_only else a for a in classes_a]
                 batch_indices_a = [ self.class_indices_a[c][next(samplers_a[c])] for c in classes_a]                                 
                 batch_indices_b = [ self.class_indices_b[c][next(samplers_b[c])] for c in classes_b]
             else:
                 classes_b = list(WeightedRandomSampler(self.class_weights_b, self.batch_size, replacement=self.use_with_replacement))
-                a_list = list(self.a_only)
+                a_list = list(self.counts_a)
                 random.shuffle(a_list)
                 classes_a = [ self._pop_leave_last(a_list) if b in self.b_only else b for b in classes_b]
                 batch_indices_a = [ self.class_indices_a[c][next(samplers_a[c])] for c in classes_a]

tmnt/distribution.py

@@ -168,7 +168,7 @@ class LogisticGaussianDistribution(BaseDistribution):
 class VonMisesDistribution(BaseDistribution):
     
     def __init__(self, enc_size, n_latent, kappa=100.0, dr=0.1, device='cpu'):
-        super(VonMisesDistribution, self).__init__(enc_size, n_latent, device)
+        super(VonMisesDistribution, self).__init__(enc_size, n_latent, device, on_simplex=False)
         self.device = device
         self.kappa = kappa
         self.kld_v = torch.tensor(VonMisesDistribution._vmf_kld(self.kappa, self.n_latent), device=device)

tmnt/estimator.py

@@ -1231,6 +1231,7 @@ class SeqBowEstimator(BaseEstimator):
         llm_base_model = get_llm_model(self.llm_model_name).to(self.device)
         model = SeqBowVED(llm_base_model, self.latent_distribution, num_classes=self.n_labels, device=self.device, 
                           vocab_size = len(self.vocabulary), use_pooling = (self.llm_model_name.startswith("sentence-transformers")),
+                          entropy_loss_coef=self.entropy_loss_coef,
                           dropout=self.classifier_dropout)
         return model
 
@@ -1583,10 +1584,10 @@ class SeqBowEstimator(BaseEstimator):
 
 class SeqBowMetricEstimator(SeqBowEstimator):
 
-    def __init__(self, *args, sdml_smoothing_factor=0.3, non_scoring_index=-1, **kwargs):
+    def __init__(self, *args, sdml_smoothing_factor=0.3, metric_loss_temp=0.1, non_scoring_index=-1, **kwargs):
         super(SeqBowMetricEstimator, self).__init__(*args, **kwargs)
         #self.loss_function = GeneralizedSDMLLoss(smoothing_parameter=sdml_smoothing_factor, x2_downweight_idx=non_scoring_index)
-        self.loss_function = MultiNegativeCrossEntropyLoss(smoothing_parameter=sdml_smoothing_factor)
+        self.loss_function = MultiNegativeCrossEntropyLoss(smoothing_parameter=sdml_smoothing_factor, metric_loss_temp=metric_loss_temp)
         self.non_scoring_index = non_scoring_index ## if >=0 this will avoid considering this label index in evaluation
 
 

tmnt/modeling.py

@@ -506,7 +506,7 @@ class BaseSeqBowVED(BaseVAE):
             input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
             return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
         else:
-            model_output.last_hidden_state[:,0,:]
+            return model_output.last_hidden_state[:,0,:]
 
     def get_ordered_terms(self):
         """
@@ -551,8 +551,8 @@ class SeqBowVED(BaseSeqBowVED):
         super(SeqBowVED, self).__init__(*args, **kwargs)
         if self.has_classifier:
             self.classifier = torch.nn.Sequential()
-            self.classifier.add_module("dr", nn.Dropout(self.dropout))
-            self.classifier.add_module("l_out", nn.Linear(self.n_latent, self.num_classes))
+            self.classifier.add_module("dr", nn.Dropout(self.dropout).to(self.device))
+            self.classifier.add_module("l_out", nn.Linear(self.n_latent, self.num_classes).to(self.device))
         
     def forward(self, input_ids, attention_mask, bow=None):  # pylint: disable=arguments-differ
         llm_output = self.llm(input_ids, attention_mask)

{tmnt-0.7.0b20230910.dist-info → tmnt-0.7.0b20230912.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tmnt
-Version: 0.7.0b20230910
+Version: 0.7.0b20230912
 Summary: Topic modeling neural toolkit
 Home-page: https://github.com/mitre/tmnt.git
 Author: The MITRE Corporation

{tmnt-0.7.0b20230910.dist-info → tmnt-0.7.0b20230912.dist-info}/RECORD

@@ -2,12 +2,12 @@ tmnt/__init__.py,sha256=EPNq1H7UMyMewWT_zTGBaC7ZouvCywX_gMX4G1dtmvw,250
 tmnt/bert_handling.py,sha256=4l78pzLjK0rbsGa3YxCsfVEndJPzaXTaj_928ZPZfSk,24677
 tmnt/common_params.py,sha256=uNWs1UuaTx6xlxbS5LailnuXhuIKCg8kCJqxes-kAGY,2547
 tmnt/configuration.py,sha256=P8PEhzVPKO5xG0FrdTLRQ60OYWigbzPY-OSx_hzQlrY,10054
-tmnt/data_loading.py,sha256=msX-l7ov0jynqxUvhPFe-5b6zc5FOlLiF1ruJXOWMrU,18485
-tmnt/distribution.py,sha256=-rKCVzpdu8P2NfBOB7QdZoFYGcCT3Q9K9x6fK0tRwew,8364
-tmnt/estimator.py,sha256=Fk0JQ0mmU9HTOrYf_q7jXsXIN1KKM2n1kRK_sinQmp4,78555
+tmnt/data_loading.py,sha256=Fnn3Pdrw16e6IR_QEPusiUfSCrHlk-3ddKeyzQW_5JE,18569
+tmnt/distribution.py,sha256=JrJe2HaF2uub0S8RxLAjSykg_AF6atvgUhNWbbaxaMo,8382
+tmnt/estimator.py,sha256=lMcuVvPIBiI1ChUpllCFR6ygWnclpcOcYA-ht5r4W3s,78680
 tmnt/eval_npmi.py,sha256=ODRDMsBgDM__iCNEX399ck7bAhl7ydvgDqmpfR7Y-q4,5048
 tmnt/inference.py,sha256=Hc0PRmUBLr9YbfqAGyw6-1BQqiwUUEtS8ehsWkr7QJk,18399
-tmnt/modeling.py,sha256=mWAlyqTYjZrkztYPrlvb24dtuvyMQ7v7mzyQNLOBs6o,32901
+tmnt/modeling.py,sha256=Zw6F3PZcaJsUZVuyk-sQIMNWMaqcjQH5cetZTmcqI7g,32940
 tmnt/selector.py,sha256=DWJlbdWKNxJmyLI9IRxCa3FDmaCilxTDzNfIf5mpBqc,9578
 tmnt/trainer.py,sha256=xaJtU_vHAPbos9q86NNdnwz7kpUF5BxGyTWbDG_NBA0,25802
 tmnt/classifier/__init__.py,sha256=1gLyJjCMHmnWdf-J4gnRs4uhbebtvCs9RgnZze1HTXY,67
@@ -30,9 +30,9 @@ tmnt/utils/ngram_helpers.py,sha256=VrIzou2oQHCLBLSWODDeikN3PYat1NqqvEeYQj_GhbA,1
 tmnt/utils/pubmed_utils.py,sha256=3sHwoun7vxb0GV-arhpXLMUbAZne0huAh9xQNy6H40E,1274
 tmnt/utils/random.py,sha256=qY75WG3peWoMh9pUyCPBEo6q8IvkF6VRjeb5CqJOBF8,327
 tmnt/utils/recalibrate.py,sha256=TmpB8An8bslICZ13UTJfIvr8VoqiSedtpHxec4n8CHk,1439
-tmnt-0.7.0b20230910.dist-info/LICENSE,sha256=qFZJrfJ7Zi4IXDiyiGVrHWic_l1h2tc36tI8Z7rK9bs,11356
-tmnt-0.7.0b20230910.dist-info/METADATA,sha256=cLvhJfcBDDppP-t0GF0IPmayQ6T6NfeHkKj153d7lCk,997
-tmnt-0.7.0b20230910.dist-info/NOTICE,sha256=p0kYIVAkReTFaGb4C-qPa7h5ztze6hGzOpjCMMbOipU,425
-tmnt-0.7.0b20230910.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-tmnt-0.7.0b20230910.dist-info/top_level.txt,sha256=RpYgUl187sXnqmiwKjZZdcDlHz2AALs6bGdUcukyd_E,5
-tmnt-0.7.0b20230910.dist-info/RECORD,,
+tmnt-0.7.0b20230912.dist-info/LICENSE,sha256=qFZJrfJ7Zi4IXDiyiGVrHWic_l1h2tc36tI8Z7rK9bs,11356
+tmnt-0.7.0b20230912.dist-info/METADATA,sha256=vx4UrdxOGQ2BPb9o-5mipGun99SVmyw0DOL4aS58gOU,997
+tmnt-0.7.0b20230912.dist-info/NOTICE,sha256=p0kYIVAkReTFaGb4C-qPa7h5ztze6hGzOpjCMMbOipU,425
+tmnt-0.7.0b20230912.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+tmnt-0.7.0b20230912.dist-info/top_level.txt,sha256=RpYgUl187sXnqmiwKjZZdcDlHz2AALs6bGdUcukyd_E,5
+tmnt-0.7.0b20230912.dist-info/RECORD,,