Update utils.py

utils.py

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