Spaces:
Sleeping
Sleeping
File size: 2,439 Bytes
5fde057 | 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 | from __future__ import annotations
from contextlib import nullcontext
from typing import Callable, Optional
from .environment import TicketmeltEnv
from .models import Action, Observation
from .prompt import observation_to_prompt, parse_action
def parse_action_from_model(text: str) -> Action:
"""Thin alias so tests can import this name explicitly."""
return parse_action(text)
def run_episode(
model,
tokenizer,
env: TicketmeltEnv,
prompt_fn: Callable[[Observation], str] = observation_to_prompt,
seed: Optional[int] = None,
device: str = "cuda",
max_new_tokens: int = 64,
temperature: float = 0.7,
) -> dict:
"""
Run one full episode with the given model.
Returns:
{
"history": [{"prompt": str, "response": str, "action": dict, "reward": float}, ...],
"final_reward": float,
"info": dict,
}
"""
try:
import torch
no_grad = torch.no_grad
except ImportError:
no_grad = nullcontext
obs = env.reset(seed=seed)
episode_history = []
reward = 0.0
info = {}
while not obs.done:
prompt = prompt_fn(obs)
messages = [{"role": "user", "content": prompt}]
try:
input_text = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
inputs = tokenizer(input_text, return_tensors="pt")
except Exception:
inputs = tokenizer(prompt, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
input_len = inputs["input_ids"].shape[1]
with no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
pad_token_id=tokenizer.eos_token_id,
)
response_ids = output_ids[0][input_len:]
response_text = tokenizer.decode(response_ids, skip_special_tokens=True)
action = parse_action(response_text)
obs, reward, done, info = env.step(action)
episode_history.append({
"prompt": prompt,
"response": response_text,
"action": action.to_dict(),
"reward": reward,
"done": done,
})
return {"history": episode_history, "final_reward": reward, "info": info}
|