Spaces:
Sleeping
Sleeping
File size: 8,403 Bytes
80cb919 e138b0e 80cb919 e138b0e 80cb919 e138b0e 80cb919 b0a3faf 80cb919 b0a3faf e138b0e 80cb919 e138b0e 80cb919 19d62ff 80cb919 19d62ff 80cb919 19d62ff 80cb919 |
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 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 |
# Round-robin rotator + paraphrasing + translation/backtranslation
import os
import logging
import requests
from typing import Optional
from google import genai
logger = logging.getLogger("llm")
if not logger.handlers:
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
logger.addHandler(handler)
# LLM parser limit text to log-out
def snip(s: str, n: int = 12) -> str:
if not isinstance(s, str): return "∅"
parts = s.strip().split()
return " ".join(parts[:n]) + (" …" if len(parts) > n else "")
class KeyRotator:
def __init__(self, env_prefix: str, max_keys: int = 5):
keys = []
for i in range(1, max_keys + 1):
v = os.getenv(f"{env_prefix}_{i}")
if v:
keys.append(v.strip())
if not keys:
logger.warning(f"[LLM] No keys found for prefix {env_prefix}_*")
self.keys = keys
self.dead = set()
self.idx = 0
def next_key(self) -> Optional[str]:
if not self.keys:
return None
for _ in range(len(self.keys)):
k = self.keys[self.idx % len(self.keys)]
self.idx += 1
if k not in self.dead:
return k
return None
def mark_bad(self, key: Optional[str]):
if key:
self.dead.add(key)
logger.warning(f"[LLM] Quarantined key (prefix hidden): {key[:6]}***")
class GeminiClient:
def __init__(self, rotator: KeyRotator, default_model: str):
self.rotator = rotator
self.default_model = default_model
def generate(self, prompt: str, model: Optional[str] = None, temperature: float = 0.2, max_output_tokens: int = 512) -> Optional[str]:
key = self.rotator.next_key()
if not key:
return None
try:
client = genai.Client(api_key=key)
# NOTE: matches your required pattern/use
res = client.models.generate_content(
model=model or self.default_model,
contents=prompt
)
text = getattr(res, "text", None)
if text:
logger.info(f"[LLM][Gemini] out={snip(text)}")
return text
except Exception as e:
logger.error(f"[LLM][Gemini] {e}")
self.rotator.mark_bad(key)
return None
class NvidiaClient:
def __init__(self, rotator: KeyRotator, default_model: str):
self.rotator = rotator
self.default_model = default_model
self.url = os.getenv("NVIDIA_API_URL", "https://integrate.api.nvidia.com/v1/chat/completions")
# Regex-based cleaning resp from quotes
def _clean_resp(self, resp: str) -> str:
if not resp: return resp
txt = resp.strip()
# Remove common boilerplate prefixes
for pat in [
r"^Here is (a|the) .*?:\s*",
r"^Paraphrased(?: version)?:\s*",
r"^Sure[,.]?\s*",
r"^Okay[,.]?\s*"
]:
import re
txt = re.sub(pat, "", txt, flags=re.I)
return txt.strip()
def generate(self, prompt: str, model: Optional[str] = None, temperature: float = 0.2, max_tokens: int = 512) -> Optional[str]:
key = self.rotator.next_key()
if not key:
return None
try:
headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
payload = {
"model": model or self.default_model,
"messages": [{"role": "user", "content": prompt}],
"temperature": temperature,
"max_tokens": max_tokens
}
r = requests.post(self.url, headers=headers, json=payload, timeout=45)
if r.status_code >= 400:
raise RuntimeError(f"HTTP {r.status_code}: {r.text[:200]}")
data = r.json()
text = data["choices"][0]["message"]["content"]
clean = self._clean_resp(text)
# Log the output here
logger.info(f"[LLM][NVIDIA] out={snip(clean)}")
return clean
except Exception as e:
logger.error(f"[LLM][NVIDIA] {e}")
self.rotator.mark_bad(key)
return None
class Paraphraser:
"""Prefers NVIDIA (cheap), falls back to Gemini EASY only. Also offers translate/backtranslate and a tiny consistency judge."""
def __init__(self, nvidia_model: str, gemini_model_easy: str, gemini_model_hard: str):
self.nv = NvidiaClient(KeyRotator("NVIDIA_API"), nvidia_model)
self.gm_easy = GeminiClient(KeyRotator("GEMINI_API"), gemini_model_easy)
# Only use GEMINI_MODEL_EASY, ignore hard model completely
self.gm_hard = None # Disabled - only use easy model
logger.info("Paraphraser initialized: NVIDIA -> GEMINI_EASY (GEMINI_HARD disabled)")
# Regex-based cleaning resp from quotes
def _clean_resp(self, resp: str) -> str:
if not resp: return resp
txt = resp.strip()
# Remove common boilerplate prefixes
for pat in [
r"^Here is (a|the) .*?:\s*",
r"^Paraphrased(?: version)?:\s*",
r"^Sure[,.]?\s*",
r"^Okay[,.]?\s*"
]:
import re
txt = re.sub(pat, "", txt, flags=re.I)
return txt.strip()
# ————— Paraphrase —————
def paraphrase(self, text: str, difficulty: str = "easy", custom_prompt: str = None) -> str:
if not text or len(text) < 12:
return text
# Use custom prompt if provided, otherwise use default
if custom_prompt:
prompt = custom_prompt
else:
prompt = (
"Paraphrase the following medical text concisely, preserve meaning and clinical terms.\n"
"Do not fabricate or remove factual claims.\n"
"Return ONLY the rewritten text, without any introduction, commentary.\n"+ text
)
# Always try NVIDIA first
out = self.nv.generate(prompt, temperature=0.1, max_tokens=min(600, max(128, len(text)//2)))
if out:
return self._clean_resp(out)
# Only fallback to GEMINI_MODEL_EASY (ignore difficulty parameter)
out = self.gm_easy.generate(prompt, max_output_tokens=min(600, max(128, len(text)//2)))
if out:
logger.info(f"[LLM][GEMINI] out={snip(self._clean_resp(out))}")
return self._clean_resp(out)
return text
# ————— Translate & Backtranslate —————
def translate(self, text: str, target_lang: str = "vi") -> Optional[str]:
if not text: return text
prompt = f"Translate to {target_lang}. Keep meaning exact, preserve medical terms:\n\n{text}"
out = self.nv.generate(prompt, temperature=0.0, max_tokens=min(800, len(text)+100))
if out: return out.strip()
return self.gm_easy.generate(prompt, max_output_tokens=min(800, len(text)+100))
def backtranslate(self, text: str, via_lang: str = "vi") -> Optional[str]:
if not text: return text
mid = self.translate(text, target_lang=via_lang)
if not mid: return None
prompt = f"Translate the following Vietnamese text back to English, preserving the exact meaning:\n\n{mid}"
out = self.nv.generate(prompt, temperature=0.0, max_tokens=min(900, len(text)+150))
if out: return out.strip()
res = self.gm_easy.generate(prompt, max_output_tokens=min(900, len(text)+150))
return res.strip() if res else None
# ————— Consistency Judge (cheap, ratio-based) —————
def consistency_check(self, user: str, output: str) -> bool:
"""Return True if 'output' appears supported by 'user' (context/question). Soft heuristic via LLM."""
prompt = (
"You are a strict medical QA validator. Given the USER input (question+context) "
"and the MODEL ANSWER, reply with exactly 'PASS' if the answer is supported and safe, "
"otherwise 'FAIL'. No extra text.\n\n"
f"USER:\n{user}\n\nANSWER:\n{output}"
)
out = self.nv.generate(prompt, temperature=0.0, max_tokens=3)
if not out:
out = self.gm_easy.generate(prompt, max_output_tokens=3)
return isinstance(out, str) and "PASS" in out.upper()
|