Spaces:

Shreya2005-2005
/

study-session-optimizer

Sleeping

Shreya2005-2005

CRITICAL: fix scoring formulas + switch to gemma model

ae0abd1 7 months ago

11.1 kB

#!/usr/bin/env python3
# inference.py
# 
# This is the main script evaluators run to test the environment.
# It spins up the server, uses an LLM to decide on actions, runs an episode,
# and prints logs in the required [START] [STEP] [END] format.
#
# Needed env vars:
#   API_BASE_URL - base URL for the LLM API (OpenAI-compatible)
#   MODEL_NAME   - which model to use
#   HF_TOKEN     - HuggingFace token (used if server is on HF Spaces)
#
# Usage:
#   python inference.py
#   python inference.py --task survive_8_hours
#   python inference.py --url http://localhost:8000

import os
import sys
import json
import time
import argparse
import threading
import subprocess
from typing import Optional

import httpx
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

# --- Config ---
API_BASE_URL = os.environ.get("API_BASE_URL", "https://openrouter.ai/api/v1")
MODEL_NAME = os.environ.get("MODEL_NAME", "google/gemma-2-9b-it:free")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
SERVER_URL = os.environ.get("SERVER_URL", "http://localhost:7860")

# how many steps before we force-stop (safety net)
MAX_STEPS = 100

# rate limit: space calls apart to avoid hitting limits
_last_llm_call = 0.0
MIN_CALL_INTERVAL = 3.0  # seconds between LLM calls


def wait_for_server(url: str, timeout: int = 30) -> bool:
    """Poll until the server is up. Returns True if it came up in time."""
    print(f"[START] Waiting for server at {url}...")
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            r = httpx.get(f"{url}/health", timeout=3)
            if r.status_code == 200:
                print(f"[START] Server is up!")
                return True
        except Exception:
            pass
        time.sleep(1)
    return False


def start_local_server() -> Optional[subprocess.Popen]:
    """Start the FastAPI server as a subprocess. Returns the process."""
    print("[START] Launching local server...")
    proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "server.app:app",
         "--host", "0.0.0.0", "--port", "7860", "--log-level", "warning"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    return proc


def _heuristic_action(obs: dict, task_id: str) -> int:
    """Smart fallback heuristic when LLM is unavailable."""
    energy = obs.get("energy", 50)
    focus = obs.get("focus", 50)
    consecutive = obs.get("consecutive_study_minutes", 0)

    if task_id == "smart_break_taker":
        # Only break when energy < 40
        if energy < 40:
            return 1 if energy > 20 else 2
        return 0  # study otherwise

    elif task_id == "survive_8_hours":
        # Conservative - preserve energy
        if energy < 25:
            return 2  # long break when very low
        elif energy < 45:
            return 1  # short break when low
        return 0  # study

    else:  # maximize_productivity
        # Study when focus is high, short break when focus drops
        if energy < 35:
            return 1  # short break to recover
        elif consecutive >= 40 or focus < 55:
            return 1  # break to restore focus
        return 0  # study to accumulate productive minutes


def get_llm_action(client: OpenAI, obs: dict, task_id: str, step_num: int) -> int:
    """
    Ask the LLM what action to take given the current observation.
    Retries up to 3 times on rate limit errors before falling back to heuristic.
    Returns 0, 1, or 2.
    """
    global _last_llm_call

    # Enforce minimum interval between calls to avoid rate limits
    now = time.time()
    wait = MIN_CALL_INTERVAL - (now - _last_llm_call)
    if wait > 0:
        time.sleep(wait)

    system_prompt = """You are controlling a virtual student in a study session simulation.
Your goal depends on the task:
- survive_8_hours: Keep energy above 0 for 8 hours (be conservative, rest often)
- maximize_productivity: Maximize focus>60 minutes in 4 hours (study when focused, short breaks when focus drops)
- smart_break_taker: Only take breaks when energy is below 40 (NEVER break when energy >= 40)
Available actions:
0 = study (5 minutes, drains energy and focus slightly)
1 = short break (10 minutes, recovers energy +15 and focus +20)
2 = long break (20 minutes, recovers energy +35 and focus +45)
Rules:
- For maximize_productivity: break when consecutive_study >= 40min OR focus < 55
- For smart_break_taker: ONLY break when energy < 40, never otherwise
- For survive_8_hours: break when energy < 40
Respond with ONLY a JSON object like: {"action": 0}
"""

    user_msg = f"""Step {step_num} | Task: {task_id}
Energy: {obs['energy']:.1f}/100 | Focus: {obs['focus']:.1f}/100
Consecutive study: {obs['consecutive_study_minutes']:.0f}min | Breaks: {obs['break_count']}
Productive minutes: {obs['productive_minutes']:.0f} | Session: {obs['session_minutes']:.0f}min
What action? Reply with only: {{"action": N}} where N is 0, 1, or 2."""

    # Try up to 3 times, with backoff on rate limit errors
    for attempt in range(3):
        try:
            _last_llm_call = time.time()
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_msg},
                ],
                max_tokens=30,
                temperature=0.1,
            )
            content = response.choices[0].message.content.strip()

            # parse JSON - try to extract action even from messy output
            try:
                parsed = json.loads(content)
            except json.JSONDecodeError:
                # try extracting number from response
                import re
                m = re.search(r'"action"\s*:\s*([012])', content)
                if m:
                    return int(m.group(1))
                # try just finding a digit
                m = re.search(r'\b([012])\b', content)
                if m:
                    return int(m.group(1))
                return _heuristic_action(obs, task_id)

            action = int(parsed.get("action", 0))
            return max(0, min(2, action))

        except Exception as ex:
            err_str = str(ex)
            is_rate_limit = "429" in err_str or "rate" in err_str.lower()

            if is_rate_limit and attempt < 2:
                wait_time = 8 * (attempt + 1)  # 8s, then 16s
                print(f"[STEP] Rate limited, waiting {wait_time}s (retry {attempt+2}/3)...")
                time.sleep(wait_time)
                continue
            else:
                print(f"[STEP] LLM failed (attempt {attempt+1}/3): using heuristic")
                return _heuristic_action(obs, task_id)

    return _heuristic_action(obs, task_id)


def run_episode(server_url: str, task_id: str, llm_client: OpenAI) -> dict:
    """
    Run a complete episode and return summary stats.
    This is where all the [STEP] logs happen.
    """
    http = httpx.Client(timeout=30.0)

    # reset the environment
    reset_resp = http.post(f"{server_url}/reset", json={"task_id": task_id})
    reset_resp.raise_for_status()
    reset_data = reset_resp.json()
    obs = reset_data["observation"]

    print(f"\n[START] Episode started | Task: {task_id}")
    print(f"[START] Initial state: energy={obs['energy']:.1f}, focus={obs['focus']:.1f}")
    print(f"[START] Using model: {MODEL_NAME} at {API_BASE_URL}")
    print("-" * 60)

    step_num = 0
    total_reward = 0.0
    done = False
    final_info = {}

    action_names = {0: "STUDY", 1: "SHORT_BREAK", 2: "LONG_BREAK"}

    while not done and step_num < MAX_STEPS:
        step_num += 1

        # get LLM decision
        action = get_llm_action(llm_client, obs, task_id, step_num)

        # execute action
        step_resp = http.post(f"{server_url}/step", json={"action": action})
        step_resp.raise_for_status()
        step_data = step_resp.json()

        new_obs = step_data["observation"]
        reward = step_data["reward"]
        done = step_data["done"]
        info = step_data.get("info", {})

        total_reward += reward
        final_info = info

        # required log format: [STEP]
        print(
            f"[STEP] {step_num:3d} | action={action_names[action]:<12} | "
            f"reward={reward:+.3f} | energy={new_obs['energy']:5.1f} | "
            f"focus={new_obs['focus']:5.1f} | productive={new_obs['productive_minutes']:5.0f}min | "
            f"done={done}"
        )

        obs = new_obs

        if done:
            break

    http.close()

    final_score = final_info.get("score", 0.5)
    print("-" * 60)
    print(f"[END] Episode finished after {step_num} steps")
    print(f"[END] Task: {task_id}")
    print(f"[END] Total reward: {total_reward:.4f}")
    print(f"[END] Final score: {final_score:.4f} (range: strictly 0 to 1 exclusive)")
    print(f"[END] Energy remaining: {obs['energy']:.1f}")
    print(f"[END] Productive minutes: {obs['productive_minutes']:.0f}")
    print(f"[END] All scores: {final_info.get('all_task_scores', {})}")

    return {
        "task_id": task_id,
        "steps": step_num,
        "total_reward": total_reward,
        "score": final_score,
        "info": final_info,
    }


def main():
    parser = argparse.ArgumentParser(description="Study Optimizer inference script")
    parser.add_argument("--task", default="maximize_productivity",
                        choices=["survive_8_hours", "maximize_productivity", "smart_break_taker"],
                        help="Task to run")
    parser.add_argument("--url", default=None,
                        help="Server URL (default: start local server)")
    parser.add_argument("--no-server", action="store_true",
                        help="Don't start local server, assume it's already running")
    args = parser.parse_args()

    server_url = args.url or SERVER_URL
    server_proc = None

    # start local server if not using remote
    if not args.no_server and args.url is None:
        server_proc = start_local_server()
        if not wait_for_server(server_url, timeout=30):
            print("[START] ERROR: Server failed to start in 30 seconds")
            if server_proc:
                server_proc.terminate()
            sys.exit(1)
    else:
        # just check it's up
        if not wait_for_server(server_url, timeout=10):
            print(f"[START] ERROR: Can't reach server at {server_url}")
            sys.exit(1)

    # set up OpenAI client
    llm_client = OpenAI(
    api_key=os.environ.get("OPENROUTER_API_KEY") or os.environ.get("OPENAI_API_KEY") or HF_TOKEN or "dummy-key",
    base_url=API_BASE_URL,
)

    try:
        results = run_episode(server_url, args.task, llm_client)
        print(f"\n[END] Done. Final score: {results['score']:.4f}")
    except KeyboardInterrupt:
        print("\n[END] Interrupted by user")
    except Exception as e:
        print(f"\n[END] ERROR: {e}")
        raise
    finally:
        if server_proc:
            print("[END] Shutting down local server...")
            server_proc.terminate()


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	# inference.py
	#
	# This is the main script evaluators run to test the environment.
	# It spins up the server, uses an LLM to decide on actions, runs an episode,
	# and prints logs in the required [START] [STEP] [END] format.
	#
	# Needed env vars:
	# API_BASE_URL - base URL for the LLM API (OpenAI-compatible)
	# MODEL_NAME - which model to use
	# HF_TOKEN - HuggingFace token (used if server is on HF Spaces)
	#
	# Usage:
	# python inference.py
	# python inference.py --task survive_8_hours
	# python inference.py --url http://localhost:8000

	import os
	import sys
	import json
	import time
	import argparse
	import threading
	import subprocess
	from typing import Optional

	import httpx
	from openai import OpenAI
	from dotenv import load_dotenv

	load_dotenv()

	# --- Config ---
	API_BASE_URL = os.environ.get("API_BASE_URL", "https://openrouter.ai/api/v1")
	MODEL_NAME = os.environ.get("MODEL_NAME", "google/gemma-2-9b-it:free")
	HF_TOKEN = os.environ.get("HF_TOKEN", "")
	SERVER_URL = os.environ.get("SERVER_URL", "http://localhost:7860")

	# how many steps before we force-stop (safety net)
	MAX_STEPS = 100

	# rate limit: space calls apart to avoid hitting limits
	_last_llm_call = 0.0
	MIN_CALL_INTERVAL = 3.0 # seconds between LLM calls


	def wait_for_server(url: str, timeout: int = 30) -> bool:
	"""Poll until the server is up. Returns True if it came up in time."""
	print(f"[START] Waiting for server at {url}...")
	deadline = time.time() + timeout
	while time.time() < deadline:
	try:
	r = httpx.get(f"{url}/health", timeout=3)
	if r.status_code == 200:
	print(f"[START] Server is up!")
	return True
	except Exception:
	pass
	time.sleep(1)
	return False


	def start_local_server() -> Optional[subprocess.Popen]:
	"""Start the FastAPI server as a subprocess. Returns the process."""
	print("[START] Launching local server...")
	proc = subprocess.Popen(
	[sys.executable, "-m", "uvicorn", "server.app:app",
	"--host", "0.0.0.0", "--port", "7860", "--log-level", "warning"],
	stdout=subprocess.PIPE,
	stderr=subprocess.PIPE,
	)
	return proc


	def _heuristic_action(obs: dict, task_id: str) -> int:
	"""Smart fallback heuristic when LLM is unavailable."""
	energy = obs.get("energy", 50)
	focus = obs.get("focus", 50)
	consecutive = obs.get("consecutive_study_minutes", 0)

	if task_id == "smart_break_taker":
	# Only break when energy < 40
	if energy < 40:
	return 1 if energy > 20 else 2
	return 0 # study otherwise

	elif task_id == "survive_8_hours":
	# Conservative - preserve energy
	if energy < 25:
	return 2 # long break when very low
	elif energy < 45:
	return 1 # short break when low
	return 0 # study

	else: # maximize_productivity
	# Study when focus is high, short break when focus drops
	if energy < 35:
	return 1 # short break to recover
	elif consecutive >= 40 or focus < 55:
	return 1 # break to restore focus
	return 0 # study to accumulate productive minutes


	def get_llm_action(client: OpenAI, obs: dict, task_id: str, step_num: int) -> int:
	"""
	Ask the LLM what action to take given the current observation.
	Retries up to 3 times on rate limit errors before falling back to heuristic.
	Returns 0, 1, or 2.
	"""
	global _last_llm_call

	# Enforce minimum interval between calls to avoid rate limits
	now = time.time()
	wait = MIN_CALL_INTERVAL - (now - _last_llm_call)
	if wait > 0:
	time.sleep(wait)

	system_prompt = """You are controlling a virtual student in a study session simulation.
	Your goal depends on the task:
	- survive_8_hours: Keep energy above 0 for 8 hours (be conservative, rest often)
	- maximize_productivity: Maximize focus>60 minutes in 4 hours (study when focused, short breaks when focus drops)
	- smart_break_taker: Only take breaks when energy is below 40 (NEVER break when energy >= 40)

	Available actions:
	0 = study (5 minutes, drains energy and focus slightly)
	1 = short break (10 minutes, recovers energy +15 and focus +20)
	2 = long break (20 minutes, recovers energy +35 and focus +45)

	Rules:
	- For maximize_productivity: break when consecutive_study >= 40min OR focus < 55
	- For smart_break_taker: ONLY break when energy < 40, never otherwise
	- For survive_8_hours: break when energy < 40

	Respond with ONLY a JSON object like: {"action": 0}
	"""

	user_msg = f"""Step {step_num} \| Task: {task_id}
	Energy: {obs['energy']:.1f}/100 \| Focus: {obs['focus']:.1f}/100
	Consecutive study: {obs['consecutive_study_minutes']:.0f}min \| Breaks: {obs['break_count']}
	Productive minutes: {obs['productive_minutes']:.0f} \| Session: {obs['session_minutes']:.0f}min

	What action? Reply with only: {{"action": N}} where N is 0, 1, or 2."""

	# Try up to 3 times, with backoff on rate limit errors
	for attempt in range(3):
	try:
	_last_llm_call = time.time()
	response = client.chat.completions.create(
	model=MODEL_NAME,
	messages=[
	{"role": "system", "content": system_prompt},
	{"role": "user", "content": user_msg},
	],
	max_tokens=30,
	temperature=0.1,
	)
	content = response.choices[0].message.content.strip()

	# parse JSON - try to extract action even from messy output
	try:
	parsed = json.loads(content)
	except json.JSONDecodeError:
	# try extracting number from response
	import re
	m = re.search(r'"action"\s:\s([012])', content)
	if m:
	return int(m.group(1))
	# try just finding a digit
	m = re.search(r'\b([012])\b', content)
	if m:
	return int(m.group(1))
	return _heuristic_action(obs, task_id)

	action = int(parsed.get("action", 0))
	return max(0, min(2, action))

	except Exception as ex:
	err_str = str(ex)
	is_rate_limit = "429" in err_str or "rate" in err_str.lower()

	if is_rate_limit and attempt < 2:
	wait_time = 8 * (attempt + 1) # 8s, then 16s
	print(f"[STEP] Rate limited, waiting {wait_time}s (retry {attempt+2}/3)...")
	time.sleep(wait_time)
	continue
	else:
	print(f"[STEP] LLM failed (attempt {attempt+1}/3): using heuristic")
	return _heuristic_action(obs, task_id)

	return _heuristic_action(obs, task_id)


	def run_episode(server_url: str, task_id: str, llm_client: OpenAI) -> dict:
	"""
	Run a complete episode and return summary stats.
	This is where all the [STEP] logs happen.
	"""
	http = httpx.Client(timeout=30.0)

	# reset the environment
	reset_resp = http.post(f"{server_url}/reset", json={"task_id": task_id})
	reset_resp.raise_for_status()
	reset_data = reset_resp.json()
	obs = reset_data["observation"]

	print(f"\n[START] Episode started \| Task: {task_id}")
	print(f"[START] Initial state: energy={obs['energy']:.1f}, focus={obs['focus']:.1f}")
	print(f"[START] Using model: {MODEL_NAME} at {API_BASE_URL}")
	print("-" * 60)

	step_num = 0
	total_reward = 0.0
	done = False
	final_info = {}

	action_names = {0: "STUDY", 1: "SHORT_BREAK", 2: "LONG_BREAK"}

	while not done and step_num < MAX_STEPS:
	step_num += 1

	# get LLM decision
	action = get_llm_action(llm_client, obs, task_id, step_num)

	# execute action
	step_resp = http.post(f"{server_url}/step", json={"action": action})
	step_resp.raise_for_status()
	step_data = step_resp.json()

	new_obs = step_data["observation"]
	reward = step_data["reward"]
	done = step_data["done"]
	info = step_data.get("info", {})

	total_reward += reward
	final_info = info

	# required log format: [STEP]
	print(
	f"[STEP] {step_num:3d} \| action={action_names[action]:<12} \| "
	f"reward={reward:+.3f} \| energy={new_obs['energy']:5.1f} \| "
	f"focus={new_obs['focus']:5.1f} \| productive={new_obs['productive_minutes']:5.0f}min \| "
	f"done={done}"
	)

	obs = new_obs

	if done:
	break

	http.close()

	final_score = final_info.get("score", 0.5)
	print("-" * 60)
	print(f"[END] Episode finished after {step_num} steps")
	print(f"[END] Task: {task_id}")
	print(f"[END] Total reward: {total_reward:.4f}")
	print(f"[END] Final score: {final_score:.4f} (range: strictly 0 to 1 exclusive)")
	print(f"[END] Energy remaining: {obs['energy']:.1f}")
	print(f"[END] Productive minutes: {obs['productive_minutes']:.0f}")
	print(f"[END] All scores: {final_info.get('all_task_scores', {})}")

	return {
	"task_id": task_id,
	"steps": step_num,
	"total_reward": total_reward,
	"score": final_score,
	"info": final_info,
	}


	def main():
	parser = argparse.ArgumentParser(description="Study Optimizer inference script")
	parser.add_argument("--task", default="maximize_productivity",
	choices=["survive_8_hours", "maximize_productivity", "smart_break_taker"],
	help="Task to run")
	parser.add_argument("--url", default=None,
	help="Server URL (default: start local server)")
	parser.add_argument("--no-server", action="store_true",
	help="Don't start local server, assume it's already running")
	args = parser.parse_args()

	server_url = args.url or SERVER_URL
	server_proc = None

	# start local server if not using remote
	if not args.no_server and args.url is None:
	server_proc = start_local_server()
	if not wait_for_server(server_url, timeout=30):
	print("[START] ERROR: Server failed to start in 30 seconds")
	if server_proc:
	server_proc.terminate()
	sys.exit(1)
	else:
	# just check it's up
	if not wait_for_server(server_url, timeout=10):
	print(f"[START] ERROR: Can't reach server at {server_url}")
	sys.exit(1)

	# set up OpenAI client
	llm_client = OpenAI(
	api_key=os.environ.get("OPENROUTER_API_KEY") or os.environ.get("OPENAI_API_KEY") or HF_TOKEN or "dummy-key",
	base_url=API_BASE_URL,
	)

	try:
	results = run_episode(server_url, args.task, llm_client)
	print(f"\n[END] Done. Final score: {results['score']:.4f}")
	except KeyboardInterrupt:
	print("\n[END] Interrupted by user")
	except Exception as e:
	print(f"\n[END] ERROR: {e}")
	raise
	finally:
	if server_proc:
	print("[END] Shutting down local server...")
	server_proc.terminate()


	if __name__ == "__main__":
	main()