Spaces:
Sleeping
Sleeping
File size: 4,764 Bytes
5fcd0ee 639b641 98b952a 639b641 98b952a 5fcd0ee 639b641 98b952a 639b641 5fcd0ee 639b641 98b952a 639b641 ebae6ab 98b952a 639b641 ebae6ab 639b641 5fcd0ee ebae6ab 5fcd0ee 639b641 ebae6ab 639b641 ebae6ab 98b952a | 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 | import os
class LLMBackend:
def __init__(self, backend: str = "hf", model_name: str = "meta-llama/Llama-2-7b-chat-hf"):
"""
backend: "groq" | "qwen" | "anthropic" | "openai" | "hf"
Default: HuggingFace Inference API — free tier, no local GPU needed.
Pipeline/client is lazy-loaded on first generate() call.
"""
self.backend = backend
self.model_name = model_name
self._pipe = None
self._client = None
if backend not in ("groq", "qwen", "anthropic", "openai", "hf"):
raise ValueError(f"Unknown backend: {backend!r}. Choose groq | qwen | anthropic | openai | hf")
def _get_pipe(self):
if self._pipe is None:
from transformers import pipeline
self._pipe = pipeline("text-generation", model=self.model_name, device_map="auto")
return self._pipe
def _get_client(self):
if self._client is None:
if self.backend == "groq":
from groq import Groq
self._client = Groq(api_key=os.environ["GROQ_API_KEY"])
elif self.backend == "anthropic":
import anthropic
self._client = anthropic.Anthropic()
elif self.backend == "openai":
from openai import OpenAI
self._client = OpenAI()
elif self.backend == "hf":
from huggingface_hub import InferenceClient
self._client = InferenceClient(token=os.environ.get("HF_TOKEN"))
return self._client
@staticmethod
def _strip_fences(text: str) -> str:
text = text.strip()
if text.startswith("```"):
newline = text.find("\n")
text = text[newline + 1:] if newline != -1 else text[3:]
if text.endswith("```"):
text = text[:-3].rstrip()
return text
def generate(self, system_prompt: str, user_prompt: str, max_tokens: int = 512, timeout_seconds: int = 30) -> str:
"""
All LLM calls must complete within timeout_seconds.
Raises TimeoutError if exceeded — caller handles gracefully.
"""
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(self._generate_inner, system_prompt, user_prompt, max_tokens)
try:
return future.result(timeout=timeout_seconds)
except concurrent.futures.TimeoutError:
raise TimeoutError(f"LLM call timed out after {timeout_seconds}s")
def _generate_inner(self, system_prompt: str, user_prompt: str, max_tokens: int) -> str:
if self.backend == "qwen":
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
]
out = self._get_pipe()(messages, max_new_tokens=max_tokens, return_full_text=False)
return self._strip_fences(out[0]["generated_text"])
elif self.backend == "groq":
resp = self._get_client().chat.completions.create(
model=self.model_name,
max_tokens=max_tokens,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
)
return self._strip_fences(resp.choices[0].message.content)
elif self.backend == "anthropic":
msg = self._get_client().messages.create(
model=self.model_name,
max_tokens=max_tokens,
system=system_prompt,
messages=[{"role": "user", "content": user_prompt}],
)
return self._strip_fences(msg.content[0].text)
elif self.backend == "openai":
resp = self._get_client().chat.completions.create(
model=self.model_name,
max_tokens=max_tokens,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
)
return self._strip_fences(resp.choices[0].message.content)
elif self.backend == "hf":
full_prompt = f"<s>[INST] {system_prompt}\n\n{user_prompt} [/INST]"
try:
response = self._get_client().text_generation(
full_prompt,
model=self.model_name,
max_new_tokens=max_tokens,
)
return self._strip_fences(response)
except Exception as e:
raise RuntimeError(f"HF Inference API error: {e}")
|