mirror of https://github.com/arc53/DocsGPT
Merge branch 'arc53:main' into main
Guruvignesh
9 months ago
committed by
GitHub
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.core.settings import settings
|
||||
import requests
|
||||
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):
|
||||
|
||||
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):
|
||||
context = messages[0]['content']
|
||||
user_question = messages[-1]['content']
|
||||
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
|
||||
|
||||
|
||||
response = requests.post(
|
||||
url=self.url,
|
||||
headers={
|
||||
"Content-Type": "application/json; charset=utf-8",
|
||||
},
|
||||
data=json.dumps({"input": prompt})
|
||||
)
|
||||
# Construct payload for endpoint
|
||||
payload = {
|
||||
"inputs": prompt,
|
||||
"stream": False,
|
||||
"parameters": {
|
||||
"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):
|
||||
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