Update utils.py

utils.py

@@ -4,11 +4,14 @@ utils.py
 
 # Standard imports
 import os
-from typing import List
+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"))
 
@@ -42,3 +45,102 @@ def get_embeddings(
     # 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})"