selection.py

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

class TfidfRetriever:
    def __init__(self, train_texts):
        self.vec = TfidfVectorizer(ngram_range=(1,2), min_df=1)
        self.X = self.vec.fit_transform(train_texts)
        self.texts = list(train_texts)

    def topk(self, query_text, k=1):
        q = self.vec.transform([query_text])
        sims = cosine_similarity(q, self.X)[0]
        idxs = np.argsort(-sims)[:k]
        return [(self.texts[i], float(sims[i])) for i in idxs]