Unverified Commit 3cadecf0 authored by Ying Sheng's avatar Ying Sheng Committed by GitHub
Browse files

Increase openai client limit (#886)

parent e90e3a50
...@@ -9,6 +9,7 @@ from dataclasses import dataclass, field ...@@ -9,6 +9,7 @@ from dataclasses import dataclass, field
from multiprocessing.pool import ThreadPool from multiprocessing.pool import ThreadPool
from typing import Any from typing import Any
import httpx
import jinja2 import jinja2
import numpy as np import numpy as np
import openai import openai
...@@ -70,6 +71,16 @@ class Eval: ...@@ -70,6 +71,16 @@ class Eval:
raise NotImplementedError() raise NotImplementedError()
class LargerHttpxClient(httpx.Client):
def __init__(self):
timeout_config = httpx.Timeout(3600)
limits = httpx.Limits(
max_keepalive_connections=3600,
max_connections=3600,
)
super().__init__(timeout=timeout_config, limits=limits)
class ChatCompletionSampler(SamplerBase): class ChatCompletionSampler(SamplerBase):
""" """
Sample from OpenAI's chat completion API Sample from OpenAI's chat completion API
...@@ -83,7 +94,7 @@ class ChatCompletionSampler(SamplerBase): ...@@ -83,7 +94,7 @@ class ChatCompletionSampler(SamplerBase):
temperature: float = 0.0, temperature: float = 0.0,
max_tokens: int = 2048, max_tokens: int = 2048,
): ):
self.client = OpenAI(base_url=base_url) self.client = OpenAI(base_url=base_url, http_client=LargerHttpxClient())
if model is None: if model is None:
model = self.client.models.list().data[0].id model = self.client.models.list().data[0].id
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment