mirror of https://github.com/arc53/DocsGPT
Merge pull request #352 from arc53/feature/aws-sagemaker-inference
sagemaker + llm creator classpull/355/head
commit
833e1836e1
@ -0,0 +1,20 @@
|
||||
from application.llm.openai import OpenAILLM, AzureOpenAILLM
|
||||
from application.llm.sagemaker import SagemakerAPILLM
|
||||
from application.llm.huggingface import HuggingFaceLLM
|
||||
|
||||
|
||||
|
||||
class LLMCreator:
|
||||
llms = {
|
||||
'openai': OpenAILLM,
|
||||
'azure_openai': AzureOpenAILLM,
|
||||
'sagemaker': SagemakerAPILLM,
|
||||
'huggingface': HuggingFaceLLM
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create_llm(cls, type, *args, **kwargs):
|
||||
llm_class = cls.llms.get(type.lower())
|
||||
if not llm_class:
|
||||
raise ValueError(f"No LLM class found for type {type}")
|
||||
return llm_class(*args, **kwargs)
|
@ -0,0 +1,27 @@
|
||||
from application.llm.base import BaseLLM
|
||||
from application.core.settings import settings
|
||||
import requests
|
||||
import json
|
||||
|
||||
class SagemakerAPILLM(BaseLLM):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.url = settings.SAGEMAKER_API_URL
|
||||
|
||||
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})
|
||||
)
|
||||
|
||||
return response.json()['answer']
|
||||
|
||||
def gen_stream(self, model, engine, messages, stream=True, **kwargs):
|
||||
raise NotImplementedError("Sagemaker does not support streaming")
|
Loading…
Reference in New Issue