Add support for Hugging Face Serverless Inference

Author: hvaara

bigcodebench/gen/util/hf_inference_request.py

@@ -0,0 +1,34 @@
+import time
+
+from huggingface_hub import InferenceClient
+from huggingface_hub.inference._generated.types import TextGenerationOutput
+
+
+def make_request(
+    client: InferenceClient,
+    message: str,
+    model: str,
+    temperature: float,
+    n: int,
+    max_new_tokens: int = 2048,
+) -> TextGenerationOutput:
+    response = client.text_generation(
+        model=model,
+        prompt=message,
+        do_sample=False,
+        max_new_tokens=max_new_tokens,
+    )
+
+    return response
+
+
+def make_auto_request(*args, **kwargs) -> TextGenerationOutput:
+    ret = None
+    while ret is None:
+        try:
+            ret = make_request(*args, **kwargs)
+        except Exception as e:
+            print("Unknown error. Waiting...")
+            print(e)
+            time.sleep(1)
+    return ret

bigcodebench/provider/__init__.py

@@ -68,6 +68,19 @@ def make_model(
             tokenizer_name=tokenizer_name,
             tokenizer_legacy=tokenizer_legacy,
         )
+    elif backend == "hf-inference":
+        from bigcodebench.provider.hf_inference import HuggingFaceInferenceDecoder
+
+        return HuggingFaceInferenceDecoder(
+            name=model,
+            subset=subset,
+            split=split,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            direct_completion=direct_completion,
+            instruction_prefix=instruction_prefix,
+            response_prefix=response_prefix,
+        )
     elif backend == "openai":
         from bigcodebench.provider.openai import OpenAIChatDecoder
 

bigcodebench/provider/hf_inference.py

@@ -0,0 +1,54 @@
+import os
+from typing import List
+from tqdm import tqdm
+
+from huggingface_hub import InferenceClient
+
+from bigcodebench.provider.base import DecoderBase
+from bigcodebench.gen.util.hf_inference_request import make_auto_request
+from bigcodebench.provider.utility import make_raw_chat_prompt
+
+
+class HuggingFaceInferenceDecoder(DecoderBase):
+    def __init__(self, name: str, **kwargs):
+        super().__init__(name, **kwargs)
+        self.client = InferenceClient(
+            provider="hf-inference", api_key=os.getenv("HF_INFERENCE_API_KEY")
+        )
+
+    def codegen(
+        self, prompts: List[str], do_sample: bool = True, num_samples: int = 200
+    ) -> List[str]:
+        if do_sample:
+            assert self.temperature > 0, "Temperature must be positive for sampling"
+
+        all_outputs = []
+
+        for prompt in tqdm(prompts):
+            outputs = []
+            message = (
+                prompt
+                if self.is_direct_completion()
+                else make_raw_chat_prompt(
+                    task_prompt=prompt,
+                    subset=self.subset,
+                    split=self.split,
+                    instruction_prefix=self.instruction_prefix,
+                    response_prefix=self.response_prefix,
+                    tokenizer=None,
+                )
+            )
+            ret = make_auto_request(
+                self.client,
+                message=message,
+                model=self.name,
+                n=num_samples,
+                temperature=self.temperature,
+                max_new_tokens=self.max_new_tokens,
+            )
+            outputs.append(ret)
+            all_outputs.append(outputs)
+        return all_outputs
+
+    def is_direct_completion(self) -> bool:
+        return self.direct_completion