token_optimiser / README.md
JayeshCC's picture
Upload folder using huggingface_hub
50a32fc verified
|
Raw
History Blame Contribute Delete
11.8 kB
metadata
title: Token Optimiser Environment
emoji: πŸ”€
colorFrom: indigo
colorTo: purple
sdk: docker
app_port: 8000
pinned: false
tags:
  - openenv
base_path: /web

πŸ”€ Prompt & Response Token Optimization Environment

An OpenEnv-compatible RL environment that trains agents to minimize LLM API token usage while preserving semantic quality β€” reducing AI inference costs at scale.


⚠️ Hackathon Status β€” Action Required

Submission #6 failed Phase 2 validation. All fixes have been applied and pushed. Someone needs to resubmit from the dashboard before 12 April 2026, 11:59 PM IST.

What was fixed (latest commits)

Commit Fix
c7fe490 Fixed openenv.yaml β€” tasks had wrong format (plain strings, no grader field)
3dde704 Each task now has its own grader function (grade_redundancy_stripping, etc.)
24eb57f Fixed fallback simulate domain responses + hard task always returns JSON

Bugs fixed in server/token_optimiser_environment.py

  1. Easy task fallback β€” format check was "brief explanation" (never matched "plain_brief"), response was about machine learning instead of photosynthesis
  2. Medium task fallback β€” response was about solar/renewable energy instead of Python vs JavaScript
  3. Hard task fallback β€” JSON only returned when compression ratio ≀ 0.6; now always returns valid JSON with all 5 required keys
  4. Keyword scorer β€” added photosynthesis-domain keywords (plants, sunlight, co2, oxygen, etc.)

Updated baseline (no HF_TOKEN)

Task Before After fix
Easy 0.657 ~0.75
Medium 0.673 ~0.49
Hard 0.137 ~0.60

Medium score may vary with the rule-based optimizer; with a real LLM it should score higher.

Checklist before resubmit

  • openenv.yaml has 3 tasks with correct grader fields
  • All 3 grader functions exist and are importable
  • Fallback responses match actual task domains
  • inference.py runs all 3 task rounds (TASK_EVAL_ROUNDS = 3)
  • Resubmit from dashboard β†’ https://openenv.scaler.com (or wherever the dashboard is)

πŸ“Œ Introduction

Large Language Model APIs charge per token. Verbose prompts and unconstrained responses waste tokens and money. This environment trains an AI agent to rewrite verbose prompts into concise, efficient versions that:

  • Use fewer input tokens
  • Guide the LLM toward shorter, correctly-formatted responses
  • Preserve the full semantic meaning and intent of the original request
  • Respect output format constraints (free text, bullet points, JSON)

The agent learns real-world prompt engineering β€” a critical skill for production LLM systems where cost efficiency matters at scale.


πŸ—οΈ Architecture

token_optimiser/
β”œβ”€β”€ inference.py                        # Hackathon evaluation script
β”œβ”€β”€ models.py                           # Pydantic data models (Action / Observation / State)
β”œβ”€β”€ client.py                           # Async WebSocket EnvClient
β”œβ”€β”€ openenv.yaml                        # OpenEnv deployment config
β”œβ”€β”€ Dockerfile                          # Root-level multi-stage build
β”œβ”€β”€ pyproject.toml                      # Package config & dependencies
β”œβ”€β”€ uv.lock                             # Locked dependencies
└── server/
    β”œβ”€β”€ app.py                          # FastAPI server (WebSocket + HTTP)
    β”œβ”€β”€ token_optimiser_environment.py  # Core RL environment logic
    └── requirements.txt                # Server dependencies

🧠 How It Works

Action Space

The agent submits a TokenOptimiserAction containing its optimized version of the original prompt:

class TokenOptimiserAction(Action):
    optimized_prompt: str   # Agent's rewritten, token-efficient prompt

Observation Space

After each step, the agent receives a TokenOptimiserObservation:

class TokenOptimiserObservation(Observation):
    llm_response: str    # Actual LLM response to the optimized prompt
    input_tokens: int    # Token count of the optimized prompt
    output_tokens: int   # Token count of the LLM response
    reward: float        # Step reward (0.0 – 1.0)

State Space

class TokenOptimiserState(State):
    original_prompt: str      # The verbose task prompt the agent must optimize
    task_difficulty: str      # "easy" | "medium" | "hard"
    task_index: int           # Index in task bank

πŸ“‹ Tasks

🟒 Easy β€” Redundancy Stripping

Original: "Could you possibly help me understand, if it's not too much trouble, what the word 'photosynthesis' means? I would really appreciate it if you could explain it to me in simple terms that are easy to understand."
Goal: Strip politeness filler and redundancy to a single direct question without formatting
Expected optimized: "What does photosynthesis mean? Be brief."
Max output: 30 tokens

🟑 Medium β€” Constraint Injection

Original: "I'm looking for information about the main differences between Python and JavaScript programming languages. Could you give me a thorough breakdown covering things like typing, use cases, performance, syntax style, and ecosystem so I can decide which one to learn first?"
Goal: Compress input AND inject format + exactly 5 bullet point counts into prompt
Expected optimized: "Compare Python and JavaScript (typing, use cases, performance, syntax, ecosystem) in exactly 5 bullet points."
Max output: 120 tokens

πŸ”΄ Hard β€” Multi-Key JSON Extraction

Original: "We need you to analyze our e-commerce platform data and provide strategic insights. Specifically: first identify which product categories are performing best by revenue, second tell us which geographic regions show the most growth potential, third identify which customer segments respond best to promotions, fourth suggest how we should allocate our Q3 marketing budget across channels, and fifth flag any market risks we should be watching. Please be thorough in your analysis and provide detailed reasoning for each point."
Goal: Compress 82-word multi-intent prompt and force structured JSON output with 5 exact required keys
Expected optimized: "Analyze e-commerce data based on revenue, growth regions, responsive segments, Q3 budget, and market risks. Output strictly as JSON with keys: top_categories, growth_regions, responsive_segments, budget_allocation, risks_watch."
Max output: 200 tokens


πŸ† Reward Function

Hybrid grading β€” combines token efficiency + LLM-as-judge semantic scoring:

Component Weight Description
Token Efficiency 0.0 – 0.40 Tokens saved vs. reference (input + output combined)
Semantic Quality 0.0 – 0.30 LLM judge rates response quality 0–10
Format Compliance 0.0 – 0.20 Response matches required format (bullets / JSON / brief)
Length Penalty βˆ’0.10 Output exceeds 2Γ— max token budget
reward = token_efficiency + (semantic_score Γ— 0.3) + format_score + length_penalty
reward = clamp(reward, 0.0, 1.0)
Score Meaning
0.0 Meaning lost, system failure, or no optimization
0.3 – 0.5 Token reduction achieved but quality degraded
0.6 – 0.8 Good balance of compression and quality
0.9 – 1.0 Optimal: minimal tokens, correct format, meaning preserved

βš™οΈ Setup & Installation

Prerequisites

  • Python 3.10+
  • uv (recommended) or pip
  • A Hugging Face account with a token that has Inference Providers permission

1. Clone and Install

git clone <your-repo-url>
cd token_optimiser

# Install with uv (recommended)
uv sync

# Or with pip
pip install -e .

2. Set Environment Variables

# Required
export HF_TOKEN="hf_your_token_here"

# Optional (these are the defaults)
export API_BASE_URL="https://router.huggingface.co/v1"
export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct"
export SERVER_URL="http://localhost:8000"

On Windows (PowerShell):

$env:HF_TOKEN = "hf_your_token_here"

HF Token Permissions: Your token must have "Make calls to Inference Providers" enabled.
Create/edit at β†’ https://huggingface.co/settings/tokens


πŸš€ Running Locally

Step 1 β€” Start the Environment Server

# Terminal 1
uv run uvicorn server.app:app --host 0.0.0.0 --port 8000

Expected output:

INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000

Step 2 β€” Run the Inference Script

# Terminal 2
uv run inference.py

Expected output:

[START] task=token_optimization env=token_optimiser model=Qwen/Qwen2.5-72B-Instruct
[STEP] step=1 action='...' reward=0.71 done=false error=null
[STEP] step=4 action='...' reward=0.70 done=false error=null
[STEP] step=7 action='...' reward=0.17 done=false error=null
[END] success=false steps=9 score=0.487 rewards=0.71,0.66,0.60,0.70,0.68,0.64,0.17,0.14,0.10

πŸ“Š Baseline Performance

Current reproducible local baseline, measured with HF_TOKEN unset and the deterministic fallback path against the three-task cycle:

Task Score
Easy 0.657
Medium 0.673
Hard 0.137
Aggregate 0.487

When HF_TOKEN is available, inference.py uses the OpenAI client against the Hugging Face router and can be rerun to regenerate model-backed scores.

Step 3 β€” (Optional) Run via Docker

# Build
docker build -t token-optimiser-env .

# Run
docker run -p 8000:8000 \
  -e HF_TOKEN=$HF_TOKEN \
  -e API_BASE_URL=$API_BASE_URL \
  token-optimiser-env

🐍 Using the Python Client

import asyncio
from token_optimiser import TokenOptimiserEnv, TokenOptimiserAction

async def main():
    async with TokenOptimiserEnv(base_url="http://localhost:8000") as env:
        # Reset β€” get the task
        result = await env.reset()
        state = await env.state()
        print(f"Task: {state.original_prompt}")
        print(f"Difficulty: {state.task_difficulty}")

        # Agent submits an optimized prompt
        result = await env.step(
            TokenOptimiserAction(optimized_prompt="Explain machine learning briefly.")
        )
        print(f"LLM Response: {result.observation.llm_response}")
        print(f"Reward: {result.reward}")
        print(f"Tokens β€” in: {result.observation.input_tokens}, out: {result.observation.output_tokens}")

asyncio.run(main())

☁️ Deployment

Deploy to Hugging Face Spaces using the OpenEnv CLI:

openenv push

πŸ”§ Environment Variables Reference

Variable Default Description
HF_TOKEN (required) Hugging Face API key with Inference Providers permission
API_BASE_URL https://router.huggingface.co/v1 LLM endpoint URL
MODEL_NAME Qwen/Qwen2.5-72B-Instruct Model used for responses and judging
SERVER_URL http://localhost:8000 Environment server URL (for inference.py)
LOCAL_IMAGE_NAME (optional) Docker image name β€” auto-spins container if set

πŸ“¦ Dependencies

  • openenv-core β‰₯ 0.2.2 β€” RL environment framework
  • openai β€” LLM API client (routed through HF)
  • fastapi + uvicorn β€” Environment server
  • pydantic v2 β€” Data model validation

See pyproject.toml for the full pinned dependency list.


πŸ“„ License

BSD License β€” see source files for details.