Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-09-20 10:35:10 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

API: add endpoint for counting tokens (#1051)

Browse source
This commit is contained in:
Tymec 2023-04-12 04:08:42 +02:00 committed by GitHub
parent 1405cd8af2
commit 832ee4323d
WARNING! Although there is a key with this ID in the database it does not verify this commit! This commit is SUSPICIOUS.
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
extensions/api/script.py
View file

@ -81,6 +81,19 @@ class Handler(BaseHTTPRequestHandler):
}]
})
self.wfile.write(response.encode('utf-8'))
elif self.path == '/api/v1/token-count':
# Not compatible with KoboldAI api
self.send_response(200)
self.send_header('Content-Type', 'application/json')
self.end_headers()
tokens = encode(body['prompt'])[0]
response = json.dumps({
'results': [{
'tokens': len(tokens)
}]
})
self.wfile.write(response.encode('utf-8'))
else:
self.send_error(404)