fix: Handle single return value. (#2396)

Co-authored-by: Lucain <lucain@huggingface.co>
huggingface · Jul 18, 2024 · 1f9dbf8 · 1f9dbf8
1 parent b5e0d76
commit 1f9dbf8
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
@@ -2251,7 +2251,12 @@ def text_generation(
         if stream:
             return _stream_text_generation_response(bytes_output, details)  # type: ignore
 
-        data = _bytes_to_dict(bytes_output)[0]  # type: ignore[arg-type]
+        data = _bytes_to_dict(bytes_output)  # type: ignore[arg-type]
+
+        # Data can be a single element (dict) or an iterable of dicts where we select the first element of.
+        if isinstance(data, list):
+            data = data[0]
+
         return TextGenerationOutput.parse_obj_as_instance(data) if details else data["generated_text"]
 
     def text_to_image(