Spaces:
Running on Zero
Running on Zero
pareshmishra
Optimize LLM latency by disabling chain-of-thought generation for direct JSON extraction
bcb01d9 | 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() | |