Spaces:
Running on Zero
Running on Zero
File size: 3,401 Bytes
10ec54c e31ce60 10ec54c e31ce60 10ec54c e31ce60 10ec54c 9e1355f e31ce60 f994262 e31ce60 f994262 10ec54c e31ce60 f994262 e31ce60 10ec54c e31ce60 10ec54c e31ce60 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """
model_loader.py
---------------
NVIDIA Build API Client for RazorShield SLM Explainer.
Operates 100% via API (Zero GPU / Zero Local Model Weight Overhead).
"""
from __future__ import annotations
import os
import logging
from typing import Any
HAS_SPACES = False
LOGGER = logging.getLogger("slm-model-loader")
class SLMModelLoader:
"""NVIDIA Build API model loader for zero-shot explanation generation."""
def __init__(
self,
model_name: str | None = None,
device: str | None = None,
max_new_tokens: int | None = None,
temperature: float | None = None,
):
if os.getenv("NVIDIA_API_KEY") and os.getenv("NVIDIA_API_KEY").strip():
if not model_name or "Qwen" in model_name:
model_name = os.getenv("NVIDIA_MODEL", "openai/gpt-oss-20b")
self.model_name = model_name or os.getenv("NVIDIA_MODEL", "openai/gpt-oss-20b")
self.max_tokens = max_new_tokens or int(os.getenv("SLM_MAX_TOKENS", "4096"))
self.temperature = temperature or float(os.getenv("SLM_TEMPERATURE", "1.0"))
self._is_loaded = False
@property
def is_loaded(self) -> bool:
"""Returns True if NVIDIA_API_KEY environment variable is configured."""
key = os.getenv("NVIDIA_API_KEY", "").strip()
return bool(key) or self._is_loaded
@is_loaded.setter
def is_loaded(self, value: bool):
self._is_loaded = value
def load_model(self) -> bool:
"""Validates NVIDIA Build API configuration."""
key = os.getenv("NVIDIA_API_KEY", "").strip()
if key:
LOGGER.info("NVIDIA_API_KEY detected. Active SLM: NVIDIA Build API (%s).", self.model_name)
self._is_loaded = True
return True
LOGGER.warning("NVIDIA_API_KEY is not set. SLM will fall back to Deterministic Explainer (0 GPU).")
self._is_loaded = False
return False
def generate(self, prompt: str) -> str:
"""Generates response using NVIDIA Build API (openai/gpt-oss-20b)."""
nvidia_key = os.getenv("NVIDIA_API_KEY", "").strip()
if not nvidia_key:
raise RuntimeError("Model is not loaded. Please set NVIDIA_API_KEY environment variable.")
try:
from openai import OpenAI
client = OpenAI(
base_url="https://integrate.api.nvidia.com/v1",
api_key=nvidia_key,
)
LOGGER.info("Executing SLM generation via NVIDIA Build API (%s)...", self.model_name)
completion = client.chat.completions.create(
model=self.model_name,
messages=[{"role": "user", "content": prompt}],
temperature=self.temperature,
top_p=1,
max_tokens=self.max_tokens,
stream=False,
)
reasoning = getattr(completion.choices[0].message, "reasoning_content", None)
if reasoning:
LOGGER.info("NVIDIA Build API reasoning trace: %s", reasoning[:200])
content = completion.choices[0].message.content
if content and content.strip():
return content.strip()
raise RuntimeError("NVIDIA Build API returned empty response content.")
except Exception as exc:
LOGGER.error("NVIDIA Build API call failed: %s", exc)
raise exc
|