add hugging face language model demo notebook

Signed-off-by: Farhan Ahmed <Farhan.Ahmed@ibm.com>
Trusted-AI · Dec 14, 2023 · 8d7b586 · 8d7b586
1 parent 49b1b41
commit 8d7b586
Show file tree

Hide file tree

Showing 2 changed files with 481 additions and 2 deletions.
diff --git a/art/estimators/language_modeling/hugging_face.py b/art/estimators/language_modeling/hugging_face.py
@@ -87,7 +87,6 @@ def __init__(
             preprocessing=preprocessing,
         )
 
-        self._model = model
         self._tokenizer = tokenizer
         self._loss = loss
         self._optimizer = optimizer
@@ -276,8 +275,10 @@ def predict(self, x: Optional[Union[str, List[str]]] = None, **kwargs) -> Dict[s
                     inputs[key] = [v_i.to(self._device) for v_i in value]
                 elif isinstance(value[0], np.ndarray):
                     inputs[key] = [torch.from_numpy(v_i).to(self._device) for v_i in value]
-                elif isinstance(value[0], (float, int)):
+                elif isinstance(value[0], list):
                     inputs[key] = torch.tensor(value).to(self._device)
+                elif isinstance(value[0], (float, int)):
+                    inputs[key] = torch.tensor([value]).to(self._device)
                 else:
                     inputs[key] = value
             else: