Spaces:
Running
Running
File size: 4,631 Bytes
bc1321c 833aed9 bc1321c 833aed9 bc1321c 833aed9 bc1321c 8441e47 bc1321c 833aed9 |
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 |
"""
utils.py
"""
# Standard imports
import os
from typing import List, Tuple
# Third party imports
import numpy as np
from google import genai
from openai import OpenAI
from sentence_transformers import SentenceTransformer
from transformers import AutoModel
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
# Maximum tokens for text-embedding-3-small
MAX_TOKENS = 8191 # We don't have access to the tokenizer for text-embedding-3-small, and just assume 1 character = 1 token here
def get_embeddings(
texts: List[str], model: str = "text-embedding-3-large"
) -> List[List[float]]:
"""
Generate embeddings for a list of texts using OpenAI API synchronously.
Args:
texts: List of strings to embed.
model: OpenAI embedding model to use (default: text-embedding-3-large).
Returns:
A list of embeddings (each embedding is a list of floats).
Raises:
Exception: If the OpenAI API call fails.
"""
# Truncate texts to max token limit
truncated_texts = [text[:MAX_TOKENS] for text in texts]
# Make the API call
response = client.embeddings.create(input=truncated_texts, model=model)
# Extract embeddings from response
embeddings = np.array([data.embedding for data in response.data])
return embeddings
MODEL_CONFIGS = {
"lionguard-2": {
"label": "LionGuard 2",
"repo_id": "govtech/lionguard-2",
"embedding_strategy": "openai",
"embedding_model": "text-embedding-3-large",
},
"lionguard-2-lite": {
"label": "LionGuard 2 Lite",
"repo_id": "govtech/lionguard-2-lite",
"embedding_strategy": "sentence_transformer",
"embedding_model": "google/embeddinggemma-300m",
},
"lionguard-2.1": {
"label": "LionGuard 2.1",
"repo_id": "govtech/lionguard-2.1",
"embedding_strategy": "gemini",
"embedding_model": "gemini-embedding-001",
},
}
DEFAULT_MODEL_KEY = "lionguard-2.1"
MODEL_CACHE = {}
EMBEDDING_MODEL_CACHE = {}
current_model_choice = DEFAULT_MODEL_KEY
GEMINI_CLIENT = None
def resolve_model_key(model_key: str = None) -> str:
key = model_key or current_model_choice
if key not in MODEL_CONFIGS:
raise ValueError(f"Unknown model selection: {key}")
return key
def load_model_instance(model_key: str):
key = resolve_model_key(model_key)
if key not in MODEL_CACHE:
repo_id = MODEL_CONFIGS[key]["repo_id"]
MODEL_CACHE[key] = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
return MODEL_CACHE[key]
def get_sentence_transformer(model_name: str):
if model_name not in EMBEDDING_MODEL_CACHE:
EMBEDDING_MODEL_CACHE[model_name] = SentenceTransformer(model_name)
return EMBEDDING_MODEL_CACHE[model_name]
def get_gemini_client():
global GEMINI_CLIENT
if GEMINI_CLIENT is None:
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
raise EnvironmentError(
"GEMINI_API_KEY environment variable is required for LionGuard 2.1."
)
GEMINI_CLIENT = genai.Client(api_key=api_key)
return GEMINI_CLIENT
def get_model_embeddings(model_key: str, texts: List[str]) -> np.ndarray:
key = resolve_model_key(model_key)
config = MODEL_CONFIGS[key]
strategy = config["embedding_strategy"]
model_name = config.get("embedding_model")
if strategy == "openai":
return get_embeddings(texts, model=model_name)
if strategy == "sentence_transformer":
embedder = get_sentence_transformer(model_name)
formatted_texts = [f"task: classification | query: {text}" for text in texts]
embeddings = embedder.encode(formatted_texts)
return np.array(embeddings)
if strategy == "gemini":
client = get_gemini_client()
result = client.models.embed_content(model=model_name, contents=texts)
return np.array([embedding.values for embedding in result.embeddings])
raise ValueError(f"Unsupported embedding strategy: {strategy}")
def predict_with_model(texts: List[str], model_key: str = None) -> Tuple[dict, str]:
key = resolve_model_key(model_key)
embeddings = get_model_embeddings(key, texts)
model = load_model_instance(key)
return model.predict(embeddings), key
def set_active_model(model_key: str) -> str:
if model_key not in MODEL_CONFIGS:
return f"⚠️ Unknown model {model_key}"
global current_model_choice
current_model_choice = model_key
load_model_instance(model_key)
label = MODEL_CONFIGS[model_key]["label"]
return f"🦁 Using {label} ({model_key})"
|