mirror of https://github.com/arc53/DocsGPT
Merge branch 'arc53:main' into main
commit
a8da4b0162
@ -0,0 +1,23 @@
|
|||||||
|
repo:
|
||||||
|
- '*'
|
||||||
|
|
||||||
|
github:
|
||||||
|
- .github/**/*
|
||||||
|
|
||||||
|
application:
|
||||||
|
- application/**/*
|
||||||
|
|
||||||
|
docs:
|
||||||
|
- docs/**/*
|
||||||
|
|
||||||
|
extensions:
|
||||||
|
- extensions/**/*
|
||||||
|
|
||||||
|
frontend:
|
||||||
|
- frontend/**/*
|
||||||
|
|
||||||
|
scripts:
|
||||||
|
- scripts/**/*
|
||||||
|
|
||||||
|
tests:
|
||||||
|
- tests/**/*
|
@ -0,0 +1,15 @@
|
|||||||
|
# https://github.com/actions/labeler
|
||||||
|
name: Pull Request Labeler
|
||||||
|
on:
|
||||||
|
- pull_request_target
|
||||||
|
jobs:
|
||||||
|
triage:
|
||||||
|
permissions:
|
||||||
|
contents: read
|
||||||
|
pull-requests: write
|
||||||
|
runs-on: ubuntu-latest
|
||||||
|
steps:
|
||||||
|
- uses: actions/labeler@v4
|
||||||
|
with:
|
||||||
|
repo-token: "${{ secrets.GITHUB_TOKEN }}"
|
||||||
|
sync-labels: true
|
@ -1,27 +1,139 @@
|
|||||||
from application.llm.base import BaseLLM
|
from application.llm.base import BaseLLM
|
||||||
from application.core.settings import settings
|
from application.core.settings import settings
|
||||||
import requests
|
|
||||||
import json
|
import json
|
||||||
|
import io
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
class LineIterator:
|
||||||
|
"""
|
||||||
|
A helper class for parsing the byte stream input.
|
||||||
|
|
||||||
|
The output of the model will be in the following format:
|
||||||
|
```
|
||||||
|
b'{"outputs": [" a"]}\n'
|
||||||
|
b'{"outputs": [" challenging"]}\n'
|
||||||
|
b'{"outputs": [" problem"]}\n'
|
||||||
|
...
|
||||||
|
```
|
||||||
|
|
||||||
|
While usually each PayloadPart event from the event stream will contain a byte array
|
||||||
|
with a full json, this is not guaranteed and some of the json objects may be split across
|
||||||
|
PayloadPart events. For example:
|
||||||
|
```
|
||||||
|
{'PayloadPart': {'Bytes': b'{"outputs": '}}
|
||||||
|
{'PayloadPart': {'Bytes': b'[" problem"]}\n'}}
|
||||||
|
```
|
||||||
|
|
||||||
|
This class accounts for this by concatenating bytes written via the 'write' function
|
||||||
|
and then exposing a method which will return lines (ending with a '\n' character) within
|
||||||
|
the buffer via the 'scan_lines' function. It maintains the position of the last read
|
||||||
|
position to ensure that previous bytes are not exposed again.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, stream):
|
||||||
|
self.byte_iterator = iter(stream)
|
||||||
|
self.buffer = io.BytesIO()
|
||||||
|
self.read_pos = 0
|
||||||
|
|
||||||
|
def __iter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __next__(self):
|
||||||
|
while True:
|
||||||
|
self.buffer.seek(self.read_pos)
|
||||||
|
line = self.buffer.readline()
|
||||||
|
if line and line[-1] == ord('\n'):
|
||||||
|
self.read_pos += len(line)
|
||||||
|
return line[:-1]
|
||||||
|
try:
|
||||||
|
chunk = next(self.byte_iterator)
|
||||||
|
except StopIteration:
|
||||||
|
if self.read_pos < self.buffer.getbuffer().nbytes:
|
||||||
|
continue
|
||||||
|
raise
|
||||||
|
if 'PayloadPart' not in chunk:
|
||||||
|
print('Unknown event type:' + chunk)
|
||||||
|
continue
|
||||||
|
self.buffer.seek(0, io.SEEK_END)
|
||||||
|
self.buffer.write(chunk['PayloadPart']['Bytes'])
|
||||||
|
|
||||||
class SagemakerAPILLM(BaseLLM):
|
class SagemakerAPILLM(BaseLLM):
|
||||||
|
|
||||||
def __init__(self, *args, **kwargs):
|
def __init__(self, *args, **kwargs):
|
||||||
self.url = settings.SAGEMAKER_API_URL
|
import boto3
|
||||||
|
runtime = boto3.client(
|
||||||
|
'runtime.sagemaker',
|
||||||
|
aws_access_key_id='xxx',
|
||||||
|
aws_secret_access_key='xxx',
|
||||||
|
region_name='us-west-2'
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
self.endpoint = settings.SAGEMAKER_ENDPOINT
|
||||||
|
self.runtime = runtime
|
||||||
|
|
||||||
|
|
||||||
def gen(self, model, engine, messages, stream=False, **kwargs):
|
def gen(self, model, engine, messages, stream=False, **kwargs):
|
||||||
context = messages[0]['content']
|
context = messages[0]['content']
|
||||||
user_question = messages[-1]['content']
|
user_question = messages[-1]['content']
|
||||||
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
|
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
|
||||||
|
|
||||||
|
|
||||||
response = requests.post(
|
# Construct payload for endpoint
|
||||||
url=self.url,
|
payload = {
|
||||||
headers={
|
"inputs": prompt,
|
||||||
"Content-Type": "application/json; charset=utf-8",
|
"stream": False,
|
||||||
},
|
"parameters": {
|
||||||
data=json.dumps({"input": prompt})
|
"do_sample": True,
|
||||||
)
|
"temperature": 0.1,
|
||||||
|
"max_new_tokens": 30,
|
||||||
|
"repetition_penalty": 1.03,
|
||||||
|
"stop": ["</s>", "###"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
body_bytes = json.dumps(payload).encode('utf-8')
|
||||||
|
|
||||||
return response.json()['answer']
|
# Invoke the endpoint
|
||||||
|
response = self.runtime.invoke_endpoint(EndpointName=self.endpoint,
|
||||||
|
ContentType='application/json',
|
||||||
|
Body=body_bytes)
|
||||||
|
result = json.loads(response['Body'].read().decode())
|
||||||
|
import sys
|
||||||
|
print(result[0]['generated_text'], file=sys.stderr)
|
||||||
|
return result[0]['generated_text'][len(prompt):]
|
||||||
|
|
||||||
def gen_stream(self, model, engine, messages, stream=True, **kwargs):
|
def gen_stream(self, model, engine, messages, stream=True, **kwargs):
|
||||||
raise NotImplementedError("Sagemaker does not support streaming")
|
context = messages[0]['content']
|
||||||
|
user_question = messages[-1]['content']
|
||||||
|
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
|
||||||
|
|
||||||
|
|
||||||
|
# Construct payload for endpoint
|
||||||
|
payload = {
|
||||||
|
"inputs": prompt,
|
||||||
|
"stream": True,
|
||||||
|
"parameters": {
|
||||||
|
"do_sample": True,
|
||||||
|
"temperature": 0.1,
|
||||||
|
"max_new_tokens": 512,
|
||||||
|
"repetition_penalty": 1.03,
|
||||||
|
"stop": ["</s>", "###"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
body_bytes = json.dumps(payload).encode('utf-8')
|
||||||
|
|
||||||
|
# Invoke the endpoint
|
||||||
|
response = self.runtime.invoke_endpoint_with_response_stream(EndpointName=self.endpoint,
|
||||||
|
ContentType='application/json',
|
||||||
|
Body=body_bytes)
|
||||||
|
#result = json.loads(response['Body'].read().decode())
|
||||||
|
event_stream = response['Body']
|
||||||
|
start_json = b'{'
|
||||||
|
for line in LineIterator(event_stream):
|
||||||
|
if line != b'' and start_json in line:
|
||||||
|
#print(line)
|
||||||
|
data = json.loads(line[line.find(start_json):].decode('utf-8'))
|
||||||
|
if data['token']['text'] not in ["</s>", "###"]:
|
||||||
|
print(data['token']['text'],end='')
|
||||||
|
yield data['token']['text']
|
@ -0,0 +1,96 @@
|
|||||||
|
# FILEPATH: /path/to/test_sagemaker.py
|
||||||
|
|
||||||
|
import json
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
from application.llm.sagemaker import SagemakerAPILLM, LineIterator
|
||||||
|
|
||||||
|
class TestSagemakerAPILLM(unittest.TestCase):
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.sagemaker = SagemakerAPILLM()
|
||||||
|
self.context = "This is the context"
|
||||||
|
self.user_question = "What is the answer?"
|
||||||
|
self.messages = [
|
||||||
|
{"content": self.context},
|
||||||
|
{"content": "Some other message"},
|
||||||
|
{"content": self.user_question}
|
||||||
|
]
|
||||||
|
self.prompt = f"### Instruction \n {self.user_question} \n ### Context \n {self.context} \n ### Answer \n"
|
||||||
|
self.payload = {
|
||||||
|
"inputs": self.prompt,
|
||||||
|
"stream": False,
|
||||||
|
"parameters": {
|
||||||
|
"do_sample": True,
|
||||||
|
"temperature": 0.1,
|
||||||
|
"max_new_tokens": 30,
|
||||||
|
"repetition_penalty": 1.03,
|
||||||
|
"stop": ["</s>", "###"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.payload_stream = {
|
||||||
|
"inputs": self.prompt,
|
||||||
|
"stream": True,
|
||||||
|
"parameters": {
|
||||||
|
"do_sample": True,
|
||||||
|
"temperature": 0.1,
|
||||||
|
"max_new_tokens": 512,
|
||||||
|
"repetition_penalty": 1.03,
|
||||||
|
"stop": ["</s>", "###"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.body_bytes = json.dumps(self.payload).encode('utf-8')
|
||||||
|
self.body_bytes_stream = json.dumps(self.payload_stream).encode('utf-8')
|
||||||
|
self.response = {
|
||||||
|
"Body": MagicMock()
|
||||||
|
}
|
||||||
|
self.result = [
|
||||||
|
{
|
||||||
|
"generated_text": "This is the generated text"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
self.response['Body'].read.return_value.decode.return_value = json.dumps(self.result)
|
||||||
|
|
||||||
|
def test_gen(self):
|
||||||
|
with patch.object(self.sagemaker.runtime, 'invoke_endpoint',
|
||||||
|
return_value=self.response) as mock_invoke_endpoint:
|
||||||
|
output = self.sagemaker.gen(None, None, self.messages)
|
||||||
|
mock_invoke_endpoint.assert_called_once_with(
|
||||||
|
EndpointName=self.sagemaker.endpoint,
|
||||||
|
ContentType='application/json',
|
||||||
|
Body=self.body_bytes
|
||||||
|
)
|
||||||
|
self.assertEqual(output,
|
||||||
|
self.result[0]['generated_text'][len(self.prompt):])
|
||||||
|
|
||||||
|
def test_gen_stream(self):
|
||||||
|
with patch.object(self.sagemaker.runtime, 'invoke_endpoint_with_response_stream',
|
||||||
|
return_value=self.response) as mock_invoke_endpoint:
|
||||||
|
output = list(self.sagemaker.gen_stream(None, None, self.messages))
|
||||||
|
mock_invoke_endpoint.assert_called_once_with(
|
||||||
|
EndpointName=self.sagemaker.endpoint,
|
||||||
|
ContentType='application/json',
|
||||||
|
Body=self.body_bytes_stream
|
||||||
|
)
|
||||||
|
self.assertEqual(output, [])
|
||||||
|
|
||||||
|
class TestLineIterator(unittest.TestCase):
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.stream = [
|
||||||
|
{'PayloadPart': {'Bytes': b'{"outputs": [" a"]}\n'}},
|
||||||
|
{'PayloadPart': {'Bytes': b'{"outputs": [" challenging"]}\n'}},
|
||||||
|
{'PayloadPart': {'Bytes': b'{"outputs": [" problem"]}\n'}}
|
||||||
|
]
|
||||||
|
self.line_iterator = LineIterator(self.stream)
|
||||||
|
|
||||||
|
def test_iter(self):
|
||||||
|
self.assertEqual(iter(self.line_iterator), self.line_iterator)
|
||||||
|
|
||||||
|
def test_next(self):
|
||||||
|
self.assertEqual(next(self.line_iterator), b'{"outputs": [" a"]}')
|
||||||
|
self.assertEqual(next(self.line_iterator), b'{"outputs": [" challenging"]}')
|
||||||
|
self.assertEqual(next(self.line_iterator), b'{"outputs": [" problem"]}')
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
Loading…
Reference in New Issue