Spaces:
Running
Running
File size: 2,466 Bytes
b68e1c5 |
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 |
import os
import time
from pathlib import Path
import joblib
import nltk
from app.chatbot import build_chat_fn
from app.config import (
BM25_PATH,
INTENT_MODEL,
NLTK_PATH,
QDRANT_API_KEY,
QDRANT_ENDPOINT,
QDRANT_MOVIE_COLLECTION_NAME,
QDRANT_TV_COLLECTION_NAME,
)
from app.llm_services import load_sentence_model
from app.retriever import get_media_retriever
from app.vectorstore import connect_qdrant
from rank_bm25 import BM25Okapi
from transformers import pipeline
start = time.time()
os.environ["TOKENIZERS_PARALLELISM"] = "false"
def load_bm25_files() -> tuple[dict[str, BM25Okapi], dict[str, int]]:
bm25_dir = Path(BM25_PATH)
try:
bm25_models = {
"movie": joblib.load(bm25_dir / "movie_bm25_model.joblib"),
"tv": joblib.load(bm25_dir / "tv_bm25_model.joblib"),
}
bm25_vocabs = {
"movie": joblib.load(bm25_dir / "movie_bm25_vocab.joblib"),
"tv": joblib.load(bm25_dir / "tv_bm25_vocab.joblib"),
}
except FileNotFoundError as e:
raise FileNotFoundError(f"Missing BM25 files: {e}")
return bm25_models, bm25_vocabs
def setup_retriever():
embed_model = load_sentence_model()
qdrant_client = connect_qdrant(endpoint=QDRANT_ENDPOINT, api_key=QDRANT_API_KEY)
nltk.data.path.append(str(NLTK_PATH))
print("✅ NLTK resources loaded")
bm25_models, bm25_vocabs = load_bm25_files()
print("✅ BM25 files loaded")
return get_media_retriever(
embed_model=embed_model,
qdrant_client=qdrant_client,
bm25_models=bm25_models,
bm25_vocabs=bm25_vocabs,
movie_collection_name=QDRANT_MOVIE_COLLECTION_NAME,
tv_collection_name=QDRANT_TV_COLLECTION_NAME,
)
def setup_intent_classifier():
print(f"🔧 Loading intent classifier from {INTENT_MODEL}")
classifier = pipeline("text-classification", model=INTENT_MODEL)
print("🔥 Warming up intent classifier...")
warmup_queries = [
"Can you recommend a feel-good movie?",
"Who directed The Godfather?",
"Do you like action films?",
]
for q in warmup_queries:
_ = classifier(q)
print("🤖 Classifier ready")
return classifier
# Initialize once at startup
retriever = setup_retriever()
intent_classifier = setup_intent_classifier()
chat_fn = build_chat_fn(retriever, intent_classifier)
print(f"🔧 Total startup time: {time.time() - start:.2f}s")
|