tplbot/initializer.py

import os
import sys
import numpy 
# Alias old numpy._core references so pickled RNG states load correctly

import pickle
import faiss
from sentence_transformers import SentenceTransformer
import google.generativeai as genai

# exposed globals
hf_client = None
index = None
docs = None
intent_clf = None
intent_le = None
generation_model = None
generation_model2 = None
generation_model3 = None
def initialize():
    global hf_client, index, docs, intent_clf, intent_le, generation_model, generation_model2, generation_model3

    hf_client = SentenceTransformer('models/all-mpnet-base-v2/all_mpnet_base_v2')
    index = faiss.read_index('tpl_rag_index_h1.faiss')
    with open('tpl_rag_docs_h1.pkl','rb') as f:
        docs = pickle.load(f)
    with open('intent_clf.pkl','rb') as f:
        intent_clf = pickle.load(f)
    with open('intent_le.pkl','rb') as f:
        intent_le = pickle.load(f)

    api_key = os.environ.get('GEMINI_API_KEY')
    if not api_key:
        raise EnvironmentError("Missing GEMINI_API_KEY")
    genai.configure(api_key=api_key)
    generation_model = genai.GenerativeModel('gemini-2.5-flash')
    generation_model2 = genai.GenerativeModel('gemini-2.0-flash-lite')
    generation_model3 = genai.GenerativeModel('gemini-2.5-flash')