Spaces:
Runtime error
Runtime error
File size: 7,465 Bytes
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 | # 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.
"""
import torch
from openenv.core.env_server.interfaces import (
Environment,
Message,
ModelTokenizer,
Transform,
)
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 _tokenize_conversation(self, conversation: list[Message]) -> torch.Tensor:
"""Tokenize a conversation with a chat-template fallback for base tokenizers."""
try:
tokens = self.tokenizer.apply_chat_template(
conversation=conversation,
tokenize=True,
return_tensors="pt", # type: ignore[arg-type]
)
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"):
try:
tokens = self.tokenizer.encode( # type: ignore[attr-defined]
fallback_text,
return_tensors="pt",
)
except TypeError:
token_ids = self.tokenizer.encode(fallback_text) # type: ignore[attr-defined]
tokens = torch.tensor([token_ids], dtype=torch.long)
else:
raise ValueError("Tokenizer must support apply_chat_template or encode")
if isinstance(tokens, torch.Tensor):
return tokens
return torch.tensor(tokens, dtype=torch.long)
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.
"""
# 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.squeeze(), 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:
# Flatten all tokens into a single 1D tensor
flattened_tokens = torch.cat(
(t.flatten() for t in self._state.history_tokens), dim=0
)
else:
flattened_tokens = torch.tensor([])
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=getattr(transformed, "tokens", torch.tensor([])),
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)
|