🤖 Deploy chat_env environment - 2025-10-21 22:49:50

Dockerfile

@@ -4,8 +4,31 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# Use the specified openenv-base image
-FROM ghcr.io/meta-pytorch/openenv-base:sha-7dd8148
+# Multi-stage build: First stage builds the base image
+FROM python:3.11-slim as base-builder
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies that all environments need
+RUN pip install --no-cache-dir \
+    fastapi>=0.104.0 \
+    "uvicorn[standard]>=0.24.0" \
+    requests>=2.25.0 \
+    wsproto>=1.0.0
+
+# Set working directory
+WORKDIR /app
+
+# Default environment variables
+ENV PYTHONPATH=/app/src
+ENV PYTHONUNBUFFERED=1
+
+# Second stage: Use the built base image and add environment-specific dependencies
+FROM base-builder
+
 # Install additional dependencies for ChatEnvironment
 RUN pip install --no-cache-dir torch transformers
 
@@ -17,6 +40,7 @@ ENV TRANSFORMERS_CACHE=/.cache
 # Pre-download the GPT-2 model to avoid permission issues during runtime
 RUN python -c "from transformers import GPT2Tokenizer; GPT2Tokenizer.from_pretrained('gpt2')"
 
+
 # Copy only what's needed for this environment
 COPY src/core/ /app/src/core/
 COPY src/envs/chat_env/ /app/src/envs/chat_env/

src/envs/chat_env/server/app.py

@@ -36,12 +36,20 @@ from .chat_environment import ChatEnvironment
 # Initialize tokenizer based on environment variable
 def get_tokenizer():
     """Get tokenizer from environment or use a mock for testing."""
-    tokenizer_name = os.environ.get("TOKENIZER_NAME", "gpt2")
+    tokenizer_name = os.environ.get("TOKENIZER_NAME", "microsoft/DialoGPT-medium")
 
     try:
         from transformers import AutoTokenizer
 
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        
+        # Set a chat template if the tokenizer doesn't have one
+        if tokenizer.chat_template is None:
+            # Use a simple chat template that works with most tokenizers
+            chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}User: {{ message['content'] }}{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}{% elif message['role'] == 'system' %}System: {{ message['content'] }}{% endif %}{% if not loop.last %}\n{% endif %}{% endfor %}"
+            tokenizer.chat_template = chat_template
+            print(f"Set custom chat template for {tokenizer_name}")
+        
         print(f"Loaded tokenizer: {tokenizer_name}")
         return tokenizer
     except ImportError:

src/envs/chat_env/server/chat_environment.py

@@ -54,9 +54,18 @@ class ChatEnvironment(Environment):
             system_message: Message = {"role": system_role, "content": system_prompt}
             self._state.history_messages.append(system_message)
             # Tokenize the system message
-            system_tokens = self.tokenizer.apply_chat_template(
-                conversation=[system_message], tokenize=True, return_tensors="pt"  # type: ignore
-            )
+            try:
+                system_tokens = self.tokenizer.apply_chat_template(
+                    conversation=[system_message], tokenize=True, return_tensors="pt"  # type: ignore
+                )
+            except (ValueError, AttributeError) as e:
+                # Fallback: tokenize the content directly if chat template fails
+                if "chat_template" in str(e).lower():
+                    role_prefix = f"{system_role}: " if system_role else ""
+                    text_to_tokenize = f"{role_prefix}{system_prompt}"
+                    system_tokens = self.tokenizer.encode(text_to_tokenize, return_tensors="pt")
+                else:
+                    raise e
             self._state.history_tokens.append(system_tokens)
 
     def reset(self) -> ChatObservation:
@@ -74,9 +83,18 @@ class ChatEnvironment(Environment):
             }
             self._state.history_messages = [system_message]
             # Tokenize the system message
-            system_tokens = self.tokenizer.apply_chat_template(
-                conversation=[system_message], tokenize=True, return_tensors="pt"  # type: ignore
-            )
+            try:
+                system_tokens = self.tokenizer.apply_chat_template(
+                    conversation=[system_message], tokenize=True, return_tensors="pt"  # type: ignore
+                )
+            except (ValueError, AttributeError) as e:
+                # Fallback: tokenize the content directly if chat template fails
+                if "chat_template" in str(e).lower():
+                    role_prefix = f"{self.system_role}: " if self.system_role else ""
+                    text_to_tokenize = f"{role_prefix}{self.system_prompt}"
+                    system_tokens = self.tokenizer.encode(text_to_tokenize, return_tensors="pt")
+                else:
+                    raise e
             self._state.history_tokens = [system_tokens]
 
         return self._create_observation()
@@ -164,9 +182,19 @@ class ChatEnvironment(Environment):
         if message["content"] is None:
             raise ValueError("Message content cannot be None")
 
-        # Tokenize the single message
-        tokens = self.tokenizer.apply_chat_template(
-            conversation=[message], tokenize=True, return_tensors="pt"  # type: ignore
-        )
+        # Try to use chat template, fallback to direct tokenization if not available
+        try:
+            tokens = self.tokenizer.apply_chat_template(
+                conversation=[message], tokenize=True, return_tensors="pt"  # type: ignore
+            )
+        except (ValueError, AttributeError) as e:
+            # Fallback: tokenize the content directly if chat template fails
+            if "chat_template" in str(e).lower():
+                # Simple fallback: just tokenize the content with role prefix
+                role_prefix = f"{message['role']}: " if message['role'] else ""
+                text_to_tokenize = f"{role_prefix}{message['content']}"
+                tokens = self.tokenizer.encode(text_to_tokenize, return_tensors="pt")
+            else:
+                raise e
 
         return ChatAction(tokens=tokens)