Vedant Sanjay Jadhav commited on
Commit
9e1355f
·
1 Parent(s): e31ce60

fix: default SLM_MODEL_NAME to openai/gpt-oss-20b for NVIDIA Build API

Browse files
Files changed (2) hide show
  1. app.py +2 -2
  2. src/explanation/model_loader.py +3 -0
app.py CHANGED
@@ -75,8 +75,8 @@ LOGGER = logging.getLogger("razorshield-app")
75
  # =============================================================================
76
 
77
  SLM_MODEL_NAME = os.getenv(
78
- "SLM_MODEL",
79
- "Qwen/Qwen2.5-0.5B-Instruct",
80
  )
81
 
82
  POLICY_MODE_DEFAULT = os.getenv(
 
75
  # =============================================================================
76
 
77
  SLM_MODEL_NAME = os.getenv(
78
+ "NVIDIA_MODEL",
79
+ os.getenv("SLM_MODEL", "openai/gpt-oss-20b"),
80
  )
81
 
82
  POLICY_MODE_DEFAULT = os.getenv(
src/explanation/model_loader.py CHANGED
@@ -25,6 +25,9 @@ class SLMModelLoader:
25
  max_new_tokens: int | None = None,
26
  temperature: float | None = None,
27
  ):
 
 
 
28
  self.model_name = model_name or os.getenv("NVIDIA_MODEL", "openai/gpt-oss-20b")
29
  self.max_tokens = max_new_tokens or int(os.getenv("SLM_MAX_TOKENS", "4096"))
30
  self.temperature = temperature or float(os.getenv("SLM_TEMPERATURE", "1.0"))
 
25
  max_new_tokens: int | None = None,
26
  temperature: float | None = None,
27
  ):
28
+ if os.getenv("NVIDIA_API_KEY") and os.getenv("NVIDIA_API_KEY").strip():
29
+ if not model_name or "Qwen" in model_name:
30
+ model_name = os.getenv("NVIDIA_MODEL", "openai/gpt-oss-20b")
31
  self.model_name = model_name or os.getenv("NVIDIA_MODEL", "openai/gpt-oss-20b")
32
  self.max_tokens = max_new_tokens or int(os.getenv("SLM_MAX_TOKENS", "4096"))
33
  self.temperature = temperature or float(os.getenv("SLM_TEMPERATURE", "1.0"))