Add functionality to ping AI2 InferD endpoints for tulu 2

fastchat/serve/api_provider.py

@@ -147,3 +147,71 @@ def palm_api_stream_iter(chat, message, temperature, top_p, max_new_tokens):
             "error_code": 0,
         }
         yield data
+
+def ai2_api_stream_iter(model_name,
+                        messages,
+                        temperature,
+                        top_p,
+                        max_new_tokens,
+                        api_key=None,
+                        api_base=None,
+                        ):
+
+    from requests import post
+    from json import loads
+
+    # get keys and needed values
+    ai2_key = api_key or os.environ.get('AI2_API_KEY')
+    api_base = api_base or "https://inferd.allen.ai/api/v1/infer"
+    model_id = "mov_01hhndc7e8k2s9x379ge8sm8xs"
+
+    # Make requests
+    gen_params = {
+        "model": model_name,
+        "prompt": messages,
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_new_tokens": max_new_tokens,
+    }
+    logger.info(f"==== request ====\n{gen_params}")
+    res = post(api_base,
+        stream=True,
+        headers={
+            "Authorization": f"Bearer {ai2_key}"
+        },
+        json={
+            "model_version_id": model_id,
+            # This input format is specific to the Tulu2 model. Other models
+            # may require different input formats. See the model's schema
+            # documentation on InferD for more information.
+            "input": {
+                "messages": messages,
+                "opts": {
+                    "max_tokens": max_new_tokens,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "logprobs": 1, # increase for more choices
+                }
+            }
+        })
+
+    if res.status_code != 200:
+        logger.error(f"unexpected response ({res.status_code}): {res.text}")
+        raise ValueError("unexpected response from InferD", res)
+
+    text = ""
+    for line in res.iter_lines():
+        if line:
+            part = loads(line)
+            if "result" in part and "output" in part["result"]:
+                for t in part["result"]["output"]["text"]:
+                    text += t
+            else:
+                logger.error(f"unexpected part: {part}")
+                raise ValueError("empty result in InferD response")
+
+            data = {
+                    "text": text,
+                    "error_code": 0,
+            }
+            yield data