Add selection.py

selection.py

@@ -0,0 +1,15 @@
+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]