fix text generation (#2982)

Author: hanouticelina

src/huggingface_hub/inference/_client.py

@@ -2400,8 +2400,8 @@ def text_generation(
         # 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"]
+        response = provider_helper.get_response(data, request_parameters)
+        return TextGenerationOutput.parse_obj_as_instance(response) if details else response["generated_text"]
 
     def text_to_image(
         self,

src/huggingface_hub/inference/_generated/_async_client.py

@@ -2456,8 +2456,8 @@ async def text_generation(
         # 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"]
+        response = provider_helper.get_response(data, request_parameters)
+        return TextGenerationOutput.parse_obj_as_instance(response) if details else response["generated_text"]
 
     async def text_to_image(
         self,

src/huggingface_hub/inference/_providers/nebius.py

@@ -14,6 +14,16 @@ class NebiusTextGenerationTask(BaseTextGenerationTask):
     def __init__(self):
         super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai")
 
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
+        output = _as_dict(response)["choices"][0]
+        return {
+            "generated_text": output["text"],
+            "details": {
+                "finish_reason": output.get("finish_reason"),
+                "seed": output.get("seed"),
+            },
+        }
+
 
 class NebiusConversationalTask(BaseConversationalTask):
     def __init__(self):

src/huggingface_hub/inference/_providers/novita.py

@@ -22,6 +22,16 @@ def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         # there is no v1/ route for novita
         return "/v3/openai/completions"
 
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
+        output = _as_dict(response)["choices"][0]
+        return {
+            "generated_text": output["text"],
+            "details": {
+                "finish_reason": output.get("finish_reason"),
+                "seed": output.get("seed"),
+            },
+        }
+
 
 class NovitaConversationalTask(BaseConversationalTask):
     def __init__(self):

src/huggingface_hub/inference/_providers/together.py

@@ -35,6 +35,16 @@ class TogetherTextGenerationTask(BaseTextGenerationTask):
     def __init__(self):
         super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
 
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
+        output = _as_dict(response)["choices"][0]
+        return {
+            "generated_text": output["text"],
+            "details": {
+                "finish_reason": output.get("finish_reason"),
+                "seed": output.get("seed"),
+            },
+        }
+
 
 class TogetherConversationalTask(BaseConversationalTask):
     def __init__(self):