From a82805d594bb49e0c51f06898b6f3585d72fbc2b Mon Sep 17 00:00:00 2001
From: bvolpato <brunocvcunha@gmail.com>
Date: Sun, 25 Oct 2026 02:02:58 -0400
Subject: [PATCH] [Python] Fix Hugging Face API error and token handling

---
 .../ml/transforms/embeddings/huggingface.py   |  5 +-
 .../transforms/embeddings/huggingface_test.py | 59 +++++++++++++++++++
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py
index c3eeccf2a0f0..edec50606158 100644
--- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py
+++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py
@@ -211,11 +211,10 @@ def load_model(self):
   def run_inference(
       self, batch, session: requests.Session, inference_args=None):
     response = session.post(
-        self._config.api_url,
-        headers=self._config.authorization_token,
-        json={
+        self._config.api_url, json={
             "inputs": batch, "options": inference_args
         })
+    response.raise_for_status()
     return response.json()
 
 
diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
index 12990cb4c2c8..111abf1350a6 100644
--- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
+++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
@@ -16,6 +16,7 @@
 
 import contextlib
 import importlib.util
+import json
 import os
 import shutil
 import sys
@@ -23,10 +24,12 @@
 import types
 import unittest
 import uuid
+from unittest import mock
 
 import jsonpickle
 import numpy as np
 import pytest
+import requests
 from parameterized import parameterized
 
 import apache_beam as beam
@@ -472,6 +475,62 @@ def test_loading_model_from_artifacts_written_by_older_versions(self):
     self.assertEqual(model.model_name, DEFAULT_MODEL_NAME)
 
 
+class InferenceAPIHandlerTest(unittest.TestCase):
+  _API_URL = 'https://example.test/embeddings'
+
+  def _handler(self, token):
+    return huggingface.InferenceAPIEmbeddings(
+        hf_token=token, columns=['text'],
+        api_url=self._API_URL).get_model_handler()
+
+  def _response(self, status_code, body):
+    response = requests.Response()
+    response.status_code = status_code
+    response.url = self._API_URL
+    response._content = json.dumps(body).encode('utf-8')
+    return response
+
+  @parameterized.expand([
+      ('explicit-token', 'explicit-token'),
+      (None, 'environment-token'),
+  ])
+  def test_run_inference_uses_session_token(self, token, expected_token):
+    embeddings = [[0.1, 0.2]]
+    with mock.patch.dict(os.environ, {'HF_TOKEN': 'environment-token'}):
+      handler = self._handler(token)
+      with handler.load_model() as session:
+        with mock.patch.object(session,
+                               'send',
+                               return_value=self._response(200,
+                                                           embeddings)) as send:
+          result = handler.run_inference(
+              ['some text'], session, inference_args={'wait_for_model': True})
+    self.assertEqual(result, embeddings)
+    request = send.call_args.args[0]
+    self.assertEqual(request.url, self._API_URL)
+    self.assertEqual(
+        request.headers['Authorization'], f'Bearer {expected_token}')
+    self.assertEqual(
+        json.loads(request.body), {
+            'inputs': ['some text'], 'options': {
+                'wait_for_model': True
+            }
+        })
+
+  @parameterized.expand([401, 429, 503])
+  def test_run_inference_raises_http_errors(self, status_code):
+    handler = self._handler('test-token')
+    response = self._response(
+        status_code, {'error': 'inference request failed'})
+    with handler.load_model() as session:
+      with mock.patch.object(session, 'send', return_value=response):
+        with self.assertRaises(requests.HTTPError) as error:
+          handler.run_inference(['some text'], session)
+    self.assertIs(error.exception.response, response)
+    self.assertIn(str(status_code), str(error.exception))
+    self.assertIn(self._API_URL, str(error.exception))
+
+
 @unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.')
 class HuggingfaceInferenceAPITest(unittest.TestCase):
   def setUp(self):