Tejasva commited on
Commit
50e4c08
·
1 Parent(s): 4b48a6b

fix: align env-var defaults with image spec requirements

Browse files
Files changed (1) hide show
  1. inference.py +13 -7
inference.py CHANGED
@@ -18,13 +18,16 @@ from env.models import Action, AgentResponse, AgentState
18
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
19
  logger = logging.getLogger(__name__)
20
 
 
 
 
21
  def setup_openai_client() -> tuple[OpenAI, str]:
22
  """Initialize OpenAI client from environment variables."""
23
 
24
  print("STEP: Initializing OpenAI client")
25
 
26
- API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
27
- MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3-8b-instruct")
28
  HF_TOKEN = os.getenv("HF_TOKEN")
29
 
30
  if not HF_TOKEN:
@@ -408,12 +411,15 @@ def train_agent(num_episodes: int = 10) -> AgentState:
408
  def main():
409
  """Main inference loop with RL training."""
410
  try:
 
 
 
411
  # Optional Docker support
412
- local_image = os.getenv("LOCAL_IMAGE_NAME")
413
- if local_image:
414
- print(f"START: inference (Docker image: {local_image})")
415
  else:
416
- print("START: inference")
 
417
 
418
  # Run RL training
419
  agent_state = train_agent(num_episodes=10)
@@ -422,7 +428,7 @@ def main():
422
  return 0
423
 
424
  except Exception as e:
425
- print(f"END: done")
426
  logger.error(f"Inference failed: {e}")
427
  return 1
428
 
 
18
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
19
  logger = logging.getLogger(__name__)
20
 
21
+ # Module-level environment variables
22
+ LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
23
+
24
  def setup_openai_client() -> tuple[OpenAI, str]:
25
  """Initialize OpenAI client from environment variables."""
26
 
27
  print("STEP: Initializing OpenAI client")
28
 
29
+ API_BASE_URL = os.getenv("API_BASE_URL", "<your-active-url>")
30
+ MODEL_NAME = os.getenv("MODEL_NAME", "<your-active-model>")
31
  HF_TOKEN = os.getenv("HF_TOKEN")
32
 
33
  if not HF_TOKEN:
 
411
  def main():
412
  """Main inference loop with RL training."""
413
  try:
414
+ # Setup client and get environment variables
415
+ client, MODEL_NAME = setup_openai_client()
416
+
417
  # Optional Docker support
418
+ if LOCAL_IMAGE_NAME:
419
+ print(f"START: inference (Docker image: {LOCAL_IMAGE_NAME})")
 
420
  else:
421
+ print("START")
422
+ print("STEP: Starting inference")
423
 
424
  # Run RL training
425
  agent_state = train_agent(num_episodes=10)
 
428
  return 0
429
 
430
  except Exception as e:
431
+ print("END")
432
  logger.error(f"Inference failed: {e}")
433
  return 1
434