Spaces:
Running
Running
File size: 7,285 Bytes
3fdcc70 eaf6e7b 3fdcc70 eaf6e7b 3fdcc70 eaf6e7b 3fdcc70 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 |
import os
from turtle import mode
import openai
import requests
from typing import (
Any,
Dict,
List,
Optional,
)
from langchain.schema import (
AIMessage,
BaseMessage,
ChatMessage,
HumanMessage,
SystemMessage,
)
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.chat_models.base import SimpleChatModel
import os
import sys
sys.path.append(os.getcwd())
from cllm.services.nlp.llms.memory import MessageMemory
from cllm.utils import timeout
class ChatOpenAI:
def __init__(
self,
model_name: str = "gpt-3.5-turbo",
temperature: float = 0.7,
model_kwargs: Dict[str, Any] = dict(),
openai_api_key: Optional[str] = None,
openai_base_url: Optional[str] = None,
) -> None:
self.model_name = model_name
self.temperature = temperature
self.model_kwargs = model_kwargs
self.api_key = os.environ.get("OPENAI_API_KEY", openai_api_key)
self.base_url = os.environ.get("OPENAI_BASE_URL", openai_base_url)
# openai.api_key = self.api_key
# openai.base_url = self.base_url
def __call__(self, messages: List[BaseMessage], **kwargs):
stream = kwargs.get("stream", False)
context = MessageMemory(messages=messages)
context.cut_memory(self.model_name)
response = self.send_message(messages=context.to_dict(), stream=stream)
return response
def get_response(self, response):
return response.choices[0].message.content
def send_message(self, messages, stream=False):
cnt = 10
while cnt > 0:
try:
result = self.get_response(
self._send_message(
model=self.model_name,
messages=messages,
temperature=self.temperature,
stream=stream,
timeout=5,
)
)
break
except Exception as e:
cnt -= 1
print(e)
result = e
return result
# @timeout(5)
def _send_message(self, *args, **kwargs):
# return self.client.chat.completions.create(*args, **kwargs)
# return openai.Completion.create(*args, **kwargs)
return openai.chat.completions.create(*args, **kwargs)
class ChatLLAMA2(SimpleChatModel):
"""Wrapper around LLAMA2
To use, you should launch you local model as web services.
"""
client: Any = None #: :meta private:
endpoint: str = "http://localhost:10051"
HUMAN_PROMPT = "user"
AI_PROMPT = "assistant"
@property
def _llm_type(self) -> str:
"""Return type of chat model."""
return "local-chat"
def _call(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
) -> str:
data = self._convert_messages_to_prompt(messages)
response = requests.post(self.endpoint, json=data)
return response.content.decode()
def _convert_one_message_to_text(self, message: BaseMessage) -> str:
if isinstance(message, ChatMessage):
message_text = {
"role": message.role.capitalize(),
"content": message.content,
}
elif isinstance(message, HumanMessage):
message_text = {"role": self.HUMAN_PROMPT, "content": message.content}
elif isinstance(message, AIMessage):
message_text = {"role": self.AI_PROMPT, "content": message.content}
elif isinstance(message, SystemMessage):
message_text = {"role": "system", "content": message.content}
else:
raise ValueError(f"Got unknown type {message}")
return message_text
def _convert_messages_to_text(self, messages: List[BaseMessage]) -> str:
"""Format a list of strings into a single string with necessary newlines.
Args:
messages (List[BaseMessage]): List of BaseMessage to combine.
Returns:
str: Combined string with necessary newlines.
"""
return [self._convert_one_message_to_text(message) for message in messages]
def _convert_messages_to_prompt(self, messages: List[BaseMessage]) -> str:
"""Format a list of messages into a full prompt for the Anthropic model
Args:
messages (List[BaseMessage]): List of BaseMessage to combine.
Returns:
str: Combined string with necessary HUMAN_PROMPT and AI_PROMPT tags.
"""
return self._convert_messages_to_text(messages)
class ChatLLAMA2(SimpleChatModel):
"""Wrapper around LLAMA2
To use, you should launch you local model as web services.
"""
client: Any = None #: :meta private:
endpoint: str = "http://localhost:10051"
HUMAN_PROMPT = "user"
AI_PROMPT = "assistant"
@property
def _llm_type(self) -> str:
"""Return type of chat model."""
return "local-chat"
def _call(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
) -> str:
data = self._convert_messages_to_prompt(messages)
response = requests.post(self.endpoint, json=data)
return response.content.decode()
def _convert_one_message_to_text(self, message: BaseMessage) -> str:
if isinstance(message, ChatMessage):
message_text = {
"role": message.role.capitalize(),
"content": message.content,
}
elif isinstance(message, HumanMessage):
message_text = {"role": self.HUMAN_PROMPT, "content": message.content}
elif isinstance(message, AIMessage):
message_text = {"role": self.AI_PROMPT, "content": message.content}
elif isinstance(message, SystemMessage):
message_text = {"role": "system", "content": message.content}
else:
raise ValueError(f"Got unknown type {message}")
return message_text
def _convert_messages_to_text(self, messages: List[BaseMessage]) -> str:
"""Format a list of strings into a single string with necessary newlines.
Args:
messages (List[BaseMessage]): List of BaseMessage to combine.
Returns:
str: Combined string with necessary newlines.
"""
return [self._convert_one_message_to_text(message) for message in messages]
def _convert_messages_to_prompt(self, messages: List[BaseMessage]) -> str:
"""Format a list of messages into a full prompt for the Anthropic model
Args:
messages (List[BaseMessage]): List of BaseMessage to combine.
Returns:
str: Combined string with necessary HUMAN_PROMPT and AI_PROMPT tags.
"""
return self._convert_messages_to_text(messages)
if __name__ == "__main__":
chat = ChatOpenAI(model_name="gpt-3.5-turbo-1106")
msg = [
SystemMessage(content="You are a helpful assistant."),
HumanMessage(content="Hello!"),
]
response = chat(msg)
print(response)
|