Increase openai client limit (#886)
This commit is contained in:
@@ -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:
|
|||||||
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):
|
|||||||
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
|
||||||
|
|||||||
Reference in New Issue
Block a user