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