NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-19 12:57:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add chat example and fix bug in managed_server

Browse source
This commit is contained in:
dmahan93 2025-10-24 23:15:56 -07:00
parent 7bf4cfbf80
commit 5d662bf1aa
2 changed files with 188 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

14
atroposlib/envs/server_handling/managed_server.py
View file

@ -10,7 +10,7 @@ This wrapper maintains a tree structure of sequences, where:
import time
import uuid
import warnings
from typing import Any, Dict, List, Optional, Union
from typing import Any, Dict, List, Optional
from openai.types.chat.chat_completion import (
ChatCompletion,
@ -50,7 +50,7 @@ class ManagedServer:
new branches. Provides proper masking for training (prompt tokens masked with -100,
logprobs set to 0.0).
Uses the clean _tokens_and_logprobs_completion_wrapper interface internally.
Uses the clean tokens_and_logprobs_completion interface internally.
"""
def __init__(
@ -300,9 +300,7 @@ class ManagedServer:
output_tokens_list,
output_logprobs_list,
finish_reasons,
) = await self.server._tokens_and_logprobs_completion_wrapper(
**completion_kwargs
)
) = await self.server.tokens_and_logprobs_completion(**completion_kwargs)
# Track each completion and build choices
n = len(output_tokens_list)
@ -360,7 +358,9 @@ class ManagedServer:
choice = Choice(
finish_reason=finish_reason,
index=i,
message=ChatCompletionMessage(content=completion_text, role="assistant"),
message=ChatCompletionMessage(
content=completion_text, role="assistant"
),
)
choices.append(choice)
@ -414,7 +414,7 @@ class ManagedServer:
output_tokens_list,
output_logprobs_list,
finish_reasons,
) = await self.server._tokens_and_logprobs_completion_wrapper(**kwargs)
) = await self.server.tokens_and_logprobs_completion(**kwargs)
# Track each completion and build choices
n = len(output_tokens_list)