Spaces:
Sleeping
Sleeping
File size: 3,791 Bytes
f8c1b4a | 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 | """Text inference client (HF Inference API, chat_completion) with graceful fallbacks."""
import json
import re
from typing import Optional, Dict, Any
try:
from huggingface_hub import InferenceClient
except ImportError:
InferenceClient = None
import config
class TextClient:
"""Wrapper for text inference via the Hugging Face Inference API."""
def __init__(self, hf_token: Optional[str] = None):
self.hf_token = hf_token or config.HF_TOKEN
self.model = config.MODEL_NAME_TEXT
self.client = None
self.available = False
if not InferenceClient:
print("[warn]huggingface_hub not installed; text generation disabled")
return
if not self.hf_token:
print("[warn]HF_TOKEN not set; text generation disabled (using fallbacks)")
return
try:
self.client = InferenceClient(token=self.hf_token, timeout=config.TIMEOUT_INFERENCE)
self.available = True
except Exception as e:
print(f"[warn]Text client initialization failed: {e}")
self.available = False
def infer(
self,
system_prompt: str,
user_prompt: str,
temperature: float = 0.65,
max_tokens: int = 512,
) -> Optional[str]:
"""Generate text via chat completion."""
if not self.available:
return None
try:
response = self.client.chat_completion(
model=self.model,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
temperature=max(temperature, 0.01),
max_tokens=max_tokens,
top_p=0.9,
)
content = response.choices[0].message.content
return content.strip() if content else None
except Exception as e:
print(f"[warn]Text inference error: {e}")
return None
def infer_json(
self,
system_prompt: str,
user_prompt: str,
temperature: float = 0.1,
max_tokens: int = 1024,
) -> Optional[Dict[str, Any]]:
"""Generate and parse a JSON response."""
raw = self.infer(system_prompt, user_prompt, temperature, max_tokens)
if not raw:
return None
# Try direct JSON parse
try:
return json.loads(raw)
except json.JSONDecodeError:
pass
# Try stripping markdown fences
try:
clean = re.sub(r"```(?:json)?\n?(.*?)\n?```", r"\1", raw, flags=re.DOTALL)
return json.loads(clean)
except json.JSONDecodeError:
pass
# Try extracting the first JSON object
try:
match = re.search(r"\{.*\}", raw, flags=re.DOTALL)
if match:
return json.loads(match.group(0))
except json.JSONDecodeError:
pass
return None
# Global client instance
_client = None
def get_client() -> TextClient:
"""Get or create the global text client."""
global _client
if _client is None:
_client = TextClient()
return _client
def infer_text(
system_prompt: str,
user_prompt: str,
temperature: float = 0.65,
max_tokens: int = 512,
) -> Optional[str]:
"""Inference wrapper function."""
return get_client().infer(system_prompt, user_prompt, temperature, max_tokens)
def infer_json(
system_prompt: str,
user_prompt: str,
temperature: float = 0.1,
max_tokens: int = 1024,
) -> Optional[Dict[str, Any]]:
"""JSON inference wrapper function."""
return get_client().infer_json(system_prompt, user_prompt, temperature, max_tokens)
|