66from typing import Any
77
88import torch
9+ import transformers
910from torch .nn .functional import softmax
1011from transformers import (
1112 AutoModelForSequenceClassification ,
1213 AutoTokenizer ,
1314 Pipeline ,
1415 TranslationPipeline ,
16+ pipeline ,
1517)
1618
1719logger = logging .Logger ("RTC_variational_pipeline" )
1820
21+ # Some methods for the
22+
23+
24+ def collate_scores (
25+ scores : list [dict [str , float ]], label_order
26+ ) -> dict [str , list | dict ]:
27+ # these need to be returned in original order
28+ # return dict for to guarantee class predictions can be recovered
29+ collated = {score ["label" ]: score ["score" ] for score in scores }
30+ return {
31+ "scores" : [collated [label ] for label in label_order ],
32+ "score_dict" : collated ,
33+ }
34+
35+
36+ def set_classifier (classifier_pars : dict , device : str ) -> transformers .Pipeline :
37+ # new helper function which given the classifier parameters sets the correct
38+ # pipeline method. This is needed because they take different kwargs
39+ # > THIS COULD BE REFACTORED BY PUTTING KWARGS IN THE CONFIG <
40+ if classifier_pars ["specific_task" ] == "zero-shot-classification" :
41+ return pipeline (
42+ task = classifier_pars ["specific_task" ],
43+ model = classifier_pars ["model" ],
44+ multi_label = True ,
45+ device = device ,
46+ ** classifier_pars .get ("kwargs" , {}),
47+ )
48+ return pipeline (
49+ task = classifier_pars ["specific_task" ],
50+ model = classifier_pars ["model" ],
51+ device = device ,
52+ ** classifier_pars .get ("kwargs" , {}),
53+ )
54+
1955
2056def set_dropout (model : torch .nn .Module , dropout_flag : bool ) -> None :
2157 """
@@ -104,7 +140,7 @@ def clean_inference(self, x):
104140 def variational_inference (self , x ):
105141 pass
106142
107- def __init__ (self , n_variational_runs = 5 , translation_batch_size = 8 ):
143+ def __init__ (self , zero_shot : bool , n_variational_runs = 5 , translation_batch_size = 8 ):
108144 # device for inference
109145 self .set_device ()
110146 debug_msg_device = f"Loading pipeline on device: { self .device }
@@ -113,7 +149,9 @@ def __init__(self, n_variational_runs=5, translation_batch_size=8):
113149 self .func_map = {
114150 "recognition" : self .recognise ,
115151 "translation" : self .translate ,
116- "classification" : self .classify_topic ,
152+ "classification" : self .classify_topic_zero_shot
153+ if zero_shot
154+ else self .classify_topic ,
117155 }
118156 # the naive outputs of the pipeline stages calculated in self.clean_inference
119157 self .naive_outputs = {
@@ -139,8 +177,10 @@ def __init__(self, n_variational_runs=5, translation_batch_size=8):
139177 self .classifier = None
140178
141179 # map pipeline names to their pipeline counterparts
142-
143- self .topic_labels = None # This should be defined in subclass if needed
180+ # to replace class descriptors, we now want class descriptors and the labels
181+ self .dataset_meta_data : dict = {
182+ None : None
183+ } # This should be defined in subclass if needed
144184
145185 def _init_pipeline_map (self ):
146186 """
@@ -193,15 +233,16 @@ def split_translate_inputs(text: str, split_key: str) -> list[str]:
193233 split_rows = split_rows [:- 1 ]
194234 return [split + split_key for split in split_rows ]
195235
196- def check_dropout (self ):
236+ @staticmethod
237+ def check_dropout (pipeline_map : transformers .Pipeline ):
197238 """
198239 Checks the existence of dropout layers in the models of the pipeline.
199240
200241 Raises:
201242 ValueError: Raised when no dropout layers are found.
202243 """
203244 logger .debug ("\n \n ------------------ Testing Dropout --------------------" )
204- for model_key , pl in self . pipeline_map .items ():
245+ for model_key , pl in pipeline_map .items ():
205246 # only test models that exist
206247 if pl is None :
207248 pipeline_none_msg_key = (
@@ -288,15 +329,41 @@ def translate(self, text: str) -> dict[str, torch.Tensor | str]:
288329 # {full translation, sentence translations, logits, semantic embeddings}
289330 return outputs
290331
291- def classify_topic (self , text : str ) -> dict [str , str ]:
332+ def classify_topic (self , text : str ) -> dict [str , list [ float ] | dict ]:
292333 """
293334 Runs the classification model
294335
295336 Returns:
296- Dictionary of classification outputs, namely the output scores.
337+ Dictionary of classification outputs, namely the output scores and
338+ label:score dictionary.
339+ """
340+ forward = self .classifier (text , top_k = None ) # type: ignore[misc]
341+ return collate_scores (forward , self .dataset_meta_data ["class_labels" ]) # type: ignore[index]
342+
343+ def classify_topic_zero_shot (self , text : str ) -> dict [str , list [float ] | dict ]:
344+ """
345+ Runs the zero-shot classification model
346+
347+ Returns:
348+ Dictionary of classification outputs, namely the output scores and
349+ label:score dictionary.
297350 """
298- forward = self .classifier (text , self .topic_labels ) # type: ignore[misc]
299- return {"scores" : forward ["scores" ]}
351+ labels = [
352+ descriptors ["en" ]
353+ for descriptors in self .dataset_meta_data ["class_descriptors" ] # type: ignore[index]
354+ ]
355+ forward = self .classifier ( # type: ignore[misc]
356+ text , labels
357+ )
358+ return collate_scores (
359+ [
360+ {"label" : label , "score" : score }
361+ for label , score in zip (
362+ forward ["labels" ], forward ["scores" ], strict = True
363+ )
364+ ],
365+ label_order = labels ,
366+ )
300367
301368 def stack_translator_sentence_metrics (
302369 self , all_sentence_metrics : list [dict [str , Any ]]
0 commit comments