Spaces:
Running on Zero
Running on Zero
File size: 5,137 Bytes
8e798a2 bcb01d9 8e798a2 bcb01d9 8e798a2 | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | import os
import json
import re
import time
import logging
from typing import Dict, Any, Optional
from openai import OpenAI, APIError, AuthenticationError, RateLimitError
from .config import settings
logger = logging.getLogger("document_intelligence.nvidia")
class NVIDIAClient:
"""Client for invoking NVIDIA LLM inference endpoints via OpenAI compatible API."""
def __init__(
self,
api_key: Optional[str] = None,
base_url: Optional[str] = None,
model: Optional[str] = None,
):
self.api_key = api_key or settings.NVIDIA_API_KEY
self.base_url = base_url or settings.NVIDIA_BASE_URL
self.model = model or settings.NVIDIA_MODEL
self._client: Optional[OpenAI] = None
def _get_client(self) -> OpenAI:
"""Lazily initialize and return the OpenAI client configured for NVIDIA."""
if not self.api_key:
raise ValueError(
"NVIDIA_API_KEY is not set. Please provide a valid NVIDIA API Key via environment variables."
)
if self._client is None:
self._client = OpenAI(
base_url=self.base_url,
api_key=self.api_key,
timeout=60.0
)
return self._client
def generate_document_intelligence(
self,
system_prompt: str,
user_prompt: str,
temperature: float = 0.2,
top_p: float = 0.95,
max_tokens: int = 4096,
) -> Dict[str, Any]:
"""
Send document extraction text and prompt to NVIDIA LLM.
Returns parsed JSON dict, reasoning content (if any), and latency.
"""
client = self._get_client()
start_time = time.time()
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
]
logger.info(
"Invoking NVIDIA model '%s' at '%s'", self.model, self.base_url
)
try:
extra_body = {}
if "nemotron" in self.model.lower():
extra_body["chat_template_kwargs"] = {"enable_thinking": False}
completion = client.chat.completions.create(
model=self.model,
messages=messages,
temperature=temperature,
top_p=top_p,
max_tokens=min(max_tokens, 2048),
extra_body=extra_body if extra_body else None,
stream=False,
)
raw_choice = completion.choices[0]
raw_content = raw_choice.message.content or ""
reasoning = getattr(raw_choice.message, "reasoning_content", None)
latency = time.time() - start_time
logger.info("NVIDIA model invocation completed in %.3f seconds", latency)
parsed_json = self._extract_json(raw_content)
return {
"parsed_data": parsed_json,
"raw_text": raw_content,
"reasoning": reasoning,
"latency_seconds": round(latency, 4),
"model": self.model,
}
except AuthenticationError as e:
logger.error("NVIDIA Authentication failed. Verify NVIDIA_API_KEY.")
raise ValueError("Authentication with NVIDIA API failed. Please check your NVIDIA_API_KEY.") from e
except RateLimitError as e:
logger.error("NVIDIA Rate Limit exceeded: %s", e)
raise RuntimeError("NVIDIA API rate limit exceeded. Please retry shortly.") from e
except APIError as e:
logger.error("NVIDIA API error: %s", e)
raise RuntimeError(f"NVIDIA API error: {e.message}") from e
except Exception as e:
logger.error("Unexpected error during NVIDIA model call: %s", e)
raise RuntimeError(f"Failed to communicate with NVIDIA LLM: {str(e)}") from e
def _extract_json(self, text: str) -> Dict[str, Any]:
"""Extract and clean valid JSON from model response text."""
cleaned = text.strip()
# Handle markdown code fences ```json ... ```
json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned)
if json_match:
candidate = json_match.group(1).strip()
try:
return json.loads(candidate)
except json.JSONDecodeError:
pass
# Handle raw json string
try:
return json.loads(cleaned)
except json.JSONDecodeError:
pass
# Try to find the first '{' and last '}'
start_idx = cleaned.find("{")
end_idx = cleaned.rfind("}")
if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
sub = cleaned[start_idx : end_idx + 1]
try:
return json.loads(sub)
except json.JSONDecodeError as e:
logger.error("Failed to parse JSON substring: %s", e)
raise ValueError(
f"Could not parse valid JSON from NVIDIA model output: {cleaned[:200]}..."
)
# Global instance
nvidia_client = NVIDIAClient()
|