wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-07-25 13:00:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

Session Close when called HuggingFace Client (#411)

Browse Source 
* [FIX] Session Close

* [NEW] Async Context Managers of Response

* [NEW] Linting
This commit is contained in:
Nil Andreu
2023-01-05 17:49:39 +01:00
committed by GitHub
parent ee50b573e1
commit 739c073328
1 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/oasst_backend/utils/hugging_face.py
+10 -8
View File
@@ -35,15 +35,17 @@ class HuggingFaceAPI:
inference: the inference we obtain from the model in HF
"""
session = aiohttp.ClientSession()
payload: Dict[str, str] = {"inputs": input}
response = await session.post(self.api_url, headers=self.headers, json=payload)
async with aiohttp.ClientSession() as session:
payload: Dict[str, str] = {"inputs": input}
# If we get a bad response
if response.status != 200:
raise OasstError("Response Error Detoxify HuggingFace", error_code=OasstErrorCode.HUGGINGFACE_API_ERROR)
async with session.post(self.api_url, headers=self.headers, json=payload) as response:
# If we get a bad response
if response.status != 200:
raise OasstError(
"Response Error Detoxify HuggingFace", error_code=OasstErrorCode.HUGGINGFACE_API_ERROR
)
# Get the response from the API call
inference = await response.json()
# Get the response from the API call
inference = await response.json()
return inference