Update models/resnet_lstm_attention/retrieval.py

models/resnet_lstm_attention/retrieval.py

@@ -57,3 +57,21 @@ class RetrievalService:
         results = [self.text_id_map[i] for i in idxs[0]]
         print(f"DEBUG: Returning results: {results}")
         return results
+
+    def text_to_text(self, text: str, top_k: int = 5):
+        with torch.no_grad():
+            emb = self.clip_model.encode_text(text).cpu().numpy()
+        emb = self._normalize(emb)
+
+        scores, idxs = self.text_index.search(emb, top_k)
+        
+        results = []
+        for j, i in enumerate(idxs[0]):
+            caption = self.text_id_map[i]  # assuming text_id_map stores the actual caption string
+            results.append({
+                "caption": caption,
+                "score": float(scores[0][j])
+            })
+        
+        print(f"DEBUG: Text-to-text results: {results}")
+        return results