94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
import os
|
|
from typing import Union, Optional, List
|
|
from loguru import logger
|
|
|
|
from langchain.callbacks import AsyncIteratorCallbackHandler
|
|
from langchain.chat_models import ChatOpenAI
|
|
from langchain.llms.base import LLM
|
|
|
|
from .llm_config import LLMConfig
|
|
# from configs.model_config import (llm_model_dict, LLM_MODEL)
|
|
|
|
|
|
class CustomLLMModel:
|
|
|
|
def __init__(self, llm: LLM):
|
|
self.llm: LLM = llm
|
|
|
|
def __call__(self, prompt: str,
|
|
stop: Optional[List[str]] = None):
|
|
return self.llm(prompt, stop)
|
|
|
|
def _call(self, prompt: str,
|
|
stop: Optional[List[str]] = None):
|
|
return self.llm(prompt, stop)
|
|
|
|
def predict(self, prompt: str,
|
|
stop: Optional[List[str]] = None):
|
|
return self.llm(prompt, stop)
|
|
|
|
def batch(self, prompts: str,
|
|
stop: Optional[List[str]] = None):
|
|
return [self.llm(prompt, stop) for prompt in prompts]
|
|
|
|
|
|
def getChatModelFromConfig(llm_config: LLMConfig, callBack: AsyncIteratorCallbackHandler = None, ) -> Union[ChatOpenAI, LLM]:
|
|
# logger.debug(f"llm type is {type(llm_config.llm)}")
|
|
if llm_config is None:
|
|
model = ChatOpenAI(
|
|
streaming=True,
|
|
verbose=True,
|
|
openai_api_key=os.environ.get("api_key"),
|
|
openai_api_base=os.environ.get("api_base_url"),
|
|
model_name=os.environ.get("LLM_MODEL", "gpt-3.5-turbo"),
|
|
temperature=os.environ.get("temperature", 0.5),
|
|
stop=os.environ.get("stop", ""),
|
|
)
|
|
return model
|
|
|
|
if llm_config and llm_config.llm and isinstance(llm_config.llm, LLM):
|
|
return CustomLLMModel(llm=llm_config.llm)
|
|
|
|
if callBack is None:
|
|
model = ChatOpenAI(
|
|
streaming=True,
|
|
verbose=True,
|
|
openai_api_key=llm_config.api_key,
|
|
openai_api_base=llm_config.api_base_url,
|
|
model_name=llm_config.model_name,
|
|
temperature=llm_config.temperature,
|
|
stop=llm_config.stop
|
|
)
|
|
else:
|
|
model = ChatOpenAI(
|
|
streaming=True,
|
|
verbose=True,
|
|
callBack=[callBack],
|
|
openai_api_key=llm_config.api_key,
|
|
openai_api_base=llm_config.api_base_url,
|
|
model_name=llm_config.model_name,
|
|
temperature=llm_config.temperature,
|
|
stop=llm_config.stop
|
|
)
|
|
|
|
return model
|
|
|
|
|
|
import json, requests
|
|
|
|
|
|
def getExtraModel():
|
|
return TestModel()
|
|
|
|
class TestModel:
|
|
|
|
def predict(self, request_body):
|
|
headers = {"Content-Type":"application/json;charset=UTF-8",
|
|
"codegpt_user":"",
|
|
"codegpt_token":""
|
|
}
|
|
xxx = requests.post(
|
|
'https://codegencore.alipay.com/api/chat/CODE_LLAMA_INT4/completion',
|
|
data=json.dumps(request_body,ensure_ascii=False).encode('utf-8'),
|
|
headers=headers)
|
|
return xxx.json()["data"] |