huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

ensure aiohttp session exists

Browse Source
This commit is contained in:
Sabidao 2024-04-21 18:10:28 +03:00 committed by GitHub
parent 3116fb5113
commit 75f954df6c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 51 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.vscode/settings.json vendored Normal file
View File

@ -0,0 +1,5 @@
{
"githubPullRequests.ignoredPullRequestBranches": [
"main"
]
}

29
clients/python/text_generation/client.py
View File

@ -1,4 +1,5 @@
import json
import aiohttp
import requests
from aiohttp import ClientSession, ClientTimeout
@ -449,11 +450,16 @@ class AsyncClient:
self.headers = headers
self.cookies = cookies
self.timeout = ClientTimeout(timeout)
self.session: ClientSession | None = None
async def __aenter__(self):
async def ensure_session(self):
if self.session or self.session.closed:
self.session = ClientSession(
headers=self.headers, cookies=self.cookies, timeout=self.timeout
)
async def __aenter__(self):
await self.ensure_session()
return self
async def __aexit__(self, exc_type, exc, tb):
@ -518,6 +524,8 @@ class AsyncClient:
The tool to use
"""
await self.ensure_session()
request = ChatRequest(
model="tgi",
messages=messages,
@ -542,10 +550,7 @@ class AsyncClient:
return self._chat_stream_response(request)
async def _chat_single_response(self, request):
async with ClientSession(
headers=self.headers, cookies=self.cookies, timeout=self.timeout
) as session:
async with session.post(
async with self.session.post(
f"{self.base_url}/v1/chat/completions", json=request.dict()
) as resp:
payload = await resp.json()
@ -643,7 +648,7 @@ class AsyncClient:
Returns:
Response: generated response
"""
await self.ensure_session()
# Validate parameters
parameters = Parameters(
best_of=best_of,
@ -667,10 +672,7 @@ class AsyncClient:
)
request = Request(inputs=prompt, stream=False, parameters=parameters)
async with ClientSession(
headers=self.headers, cookies=self.cookies, timeout=self.timeout
) as session:
async with session.post(self.base_url, json=request.dict()) as resp:
async with self.session.post(self.base_url, json=request.dict()) as resp:
payload = await resp.json()
if resp.status != 200:
@ -743,6 +745,8 @@ class AsyncClient:
AsyncIterator[StreamResponse]: stream of generated tokens
"""
# Validate parameters
await self.ensure_session()
parameters = Parameters(
best_of=None,
details=True,
@ -765,10 +769,7 @@ class AsyncClient:
)
request = Request(inputs=prompt, stream=True, parameters=parameters)
async with ClientSession(
headers=self.headers, cookies=self.cookies, timeout=self.timeout
) as session:
async with session.post(self.base_url, json=request.dict()) as resp:
async with self.session.post(self.base_url, json=request.dict()) as resp:
if resp.status != 200:
raise parse_error(resp.status, await resp.json())