Spaces:
Running
Running
File size: 7,837 Bytes
25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 6183f0f 25bcc11 | 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 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""
Chat Environment Implementation.
A chat-based environment for LLMs, designed as a blank canvas for conversation and RL.
"""
from openenv.core.env_server.interfaces import (
Environment,
Message,
ModelTokenizer,
Transform,
)
# Support both in-repo and standalone imports
try:
# In-repo imports (when running from OpenEnv repository)
from ..models import ChatAction, ChatObservation, ChatState
except ImportError as e:
if "relative import" not in str(e) and "no known parent package" not in str(e):
raise
# Standalone imports (when running via uvicorn server.app:app)
from models import ChatAction, ChatObservation, ChatState
class ChatEnvironment(Environment):
"""A chat-based environment for LLMs, designed as a blank canvas for conversation and RL.
This environment is designed to work with language models. It provides the fundamental structure
for managing conversation state but is intentionally minimal to allow maximum flexibility.
The environment owns the tokenizer and is responsible for managing both message history and tokens.
Actions contain only tokens that interface directly with models.
Args:
tokenizer: A tokenizer that will be used to tokenize the conversation
system_prompt: An optional system prompt string to use during reset calls (optional)
system_role: The role of the system (at reset time). Defaults to "system"
transform: Optional transform to apply to observations
"""
def __init__(
self,
tokenizer: ModelTokenizer,
system_prompt: str | None = None,
system_role: str = "system",
transform: Transform | None = None,
):
super().__init__(transform=transform)
if not hasattr(tokenizer, "apply_chat_template") and not hasattr(
tokenizer, "encode"
):
raise ValueError(
"Tokenizer must have 'apply_chat_template' or 'encode' method"
)
self.tokenizer = tokenizer
self.system_prompt = system_prompt
self.system_role = system_role
self._state = ChatState()
if system_prompt:
system_message: Message = {"role": system_role, "content": system_prompt}
self._state.history_messages.append(system_message)
system_tokens = self._tokenize_conversation([system_message])
self._state.history_tokens.append(system_tokens)
def _coerce_tokens(self, tokens) -> list[int]:
"""Normalize tokenizer outputs into a flat list of ints."""
if hasattr(tokens, "tolist") and callable(tokens.tolist):
tokens = tokens.tolist()
if isinstance(tokens, tuple):
tokens = list(tokens)
if isinstance(tokens, list):
flattened: list[int] = []
for token in tokens:
flattened.extend(self._coerce_tokens(token))
return flattened
return [int(tokens)]
def _tokenize_conversation(self, conversation: list[Message]) -> list[int]:
"""Tokenize a conversation with a chat-template fallback for base tokenizers."""
try:
tokens = self.tokenizer.apply_chat_template(conversation=conversation, tokenize=True)
except Exception:
# Some tokenizers (e.g. gpt2) do not define `chat_template`.
fallback_text = "".join(
f"{m['role']}: {m['content']}\n" for m in conversation
)
if hasattr(self.tokenizer, "encode"):
tokens = self.tokenizer.encode(fallback_text) # type: ignore[attr-defined]
else:
raise ValueError("Tokenizer must support apply_chat_template or encode")
return self._coerce_tokens(tokens)
def reset(self) -> ChatObservation:
"""Reset the environment to initial state.
Returns:
ChatObservation: Initial observation with system prompt (if any)
"""
self._state.history_messages = []
self._state.history_tokens = []
if self.system_prompt:
system_message: Message = {
"role": self.system_role,
"content": self.system_prompt,
}
self._state.history_messages = [system_message]
system_tokens = self._tokenize_conversation([system_message])
self._state.history_tokens = [system_tokens]
return self._create_observation()
def step(self, action: ChatAction) -> ChatObservation: # type: ignore[override]
"""Take a step in the environment by adding tokens to the chat history.
Args:
action: A ChatAction object containing tokens.
Returns:
ChatObservation: The updated observation with the new tokens added.
"""
action_tokens = [int(token) for token in action.tokens]
# Store the tokens directly from the action
self._state.history_tokens.append(action_tokens)
# Decode tokens to text and add as a message to history
decoded_text = self.tokenizer.decode(action_tokens, skip_special_tokens=True)
assistant_message: Message = {"role": "assistant", "content": decoded_text}
self._state.history_messages.append(assistant_message)
return self._create_observation()
def _create_observation(self) -> ChatObservation:
"""Create a ChatObservation from the current state.
Returns both the message history and the tokens flattened as a single tensor
ready to be used by models.
Returns:
ChatObservation: Observation with messages and flattened tokens
"""
if self._state.history_tokens:
flattened_tokens = [
token
for token_list in self._state.history_tokens
for token in token_list
]
else:
flattened_tokens = []
observation = ChatObservation(
messages=self._state.history_messages.copy(), # Copy to prevent external mutation
tokens=flattened_tokens,
)
transformed = self._apply_transform(observation)
if isinstance(transformed, ChatObservation):
return transformed
else:
# If transform returns base Observation, convert back to ChatObservation
return ChatObservation(
messages=getattr(transformed, "messages", []),
tokens=self._coerce_tokens(getattr(transformed, "tokens", [])),
done=transformed.done,
reward=transformed.reward,
)
@property
def state(self) -> ChatState:
"""Get the current state of the environment.
Returns:
ChatState: The current state.
"""
return self._state
def message_to_action(self, message: Message) -> ChatAction:
"""Convert a message dictionary to a ChatAction with tokens.
Args:
message: Dictionary with 'role' and 'content' keys
Returns:
ChatAction: A new ChatAction instance with tokenized content
Raises:
ValueError: If required keys are missing
"""
if "role" not in message:
raise ValueError("Message must contain a 'role' key")
if "content" not in message:
raise ValueError("Message must contain a 'content' key")
if message["content"] is None:
raise ValueError("Message content cannot be None")
tokens = self._tokenize_conversation([message])
return ChatAction(tokens=tokens)
|