Upload folder using huggingface_hub

inference.py

@@ -20,6 +20,10 @@ import textwrap
 from typing import List, Optional
 
 from openai import OpenAI
+try:
+    from huggingface_hub import HfFolder
+except Exception:  # pragma: no cover
+    HfFolder = None
 
 from token_optimiser import TokenOptimiserEnv, TokenOptimiserAction
 
@@ -28,10 +32,9 @@ from token_optimiser import TokenOptimiserEnv, TokenOptimiserAction
 # ---------------------------------------------------------------------------
 API_BASE_URL: str = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
 MODEL_NAME: str = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
-HF_TOKEN: Optional[str] = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
 SERVER_URL: str = os.getenv("SERVER_URL", "http://localhost:8000")
 LOCAL_IMAGE_NAME: Optional[str] = os.getenv("LOCAL_IMAGE_NAME")
-
+HF_TOKEN: Optional[str] = os.getenv("HF_TOKEN")
 TASK_NAME: str = "token_optimization"
 BENCHMARK: str = "token_optimiser"
 MAX_STEPS: int = 5
@@ -39,6 +42,28 @@ TEMPERATURE: float = 0.3
 MAX_TOKENS: int = 200
 SUCCESS_THRESHOLD: float = 0.6
 
+
+def _resolve_hf_token() -> Optional[str]:
+    """
+    Resolve API token in this order:
+    1) HF_TOKEN env var
+    2) API_KEY env var
+    3) huggingface-cli cached login token
+    """
+    token = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+    if token:
+        return token
+
+    if HfFolder is not None:
+        try:
+            return HfFolder.get_token()
+        except Exception:
+            return None
+    return None
+
+
+HF_TOKEN: Optional[str] = _resolve_hf_token()
+
 SYSTEM_PROMPT = textwrap.dedent("""
     You are a prompt optimization expert. Rewrite the given prompt to:
     1. Use the fewest possible tokens (concise language, no filler words)
@@ -257,4 +282,4 @@ async def main() -> None:
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main())

server/app.py

@@ -54,7 +54,7 @@ app = create_app(
 )
 
 
-def main(host: str = "0.0.0.0", port: int = 8000):
+def main() -> None:
     """
     Entry point for direct execution via uv run or python -m.
 
@@ -63,23 +63,20 @@ def main(host: str = "0.0.0.0", port: int = 8000):
         uv run --project . server --port 8001
         python -m token_optimiser.server.app
 
-    Args:
-        host: Host address to bind to (default: "0.0.0.0")
-        port: Port number to listen on (default: 8000)
-
     For production deployments, consider using uvicorn directly with
     multiple workers:
         uvicorn token_optimiser.server.app:app --workers 4
     """
-    import uvicorn
-
-    uvicorn.run(app, host=host, port=port)
-
-
-if __name__ == "__main__":
     import argparse
+    import uvicorn
 
     parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int, default=8000)
     args = parser.parse_args()
-    main(port=args.port)
+
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()