Add search.py

search.py

@@ -0,0 +1,44 @@
+import numpy as np
+from sentence_transformers import SentenceTransformer, util
+import pandas as pd
+import os
+
+MODEL_DIR = "ai_model"
+
+model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
+df = pd.read_csv(f"{MODEL_DIR}/search_data.csv")
+embeddings = np.load(f"{MODEL_DIR}/embeddings.npy")
+
+def get_confidence_label(score):
+    if score >= 0.7:
+        return "high"
+    elif score >= 0.5:
+        return "medium"
+    return "low"
+
+def search(query, top_k=5, threshold=0.3):
+    """Semantic search with confidence scores"""
+    q_emb = model.encode([query.lower()], normalize_embeddings=True)[0]
+    
+    scores = util.cos_sim(q_emb, embeddings)[0].numpy()
+    
+    valid_idx = np.where(scores >= threshold)[0]
+    if len(valid_idx) == 0:
+        return pd.DataFrame(columns=["rank", "somali", "english", "italian", "domain", "similarity_score", "confidence_label"])
+    
+    top_idx = valid_idx[np.argsort(scores[valid_idx])[-top_k:][::-1]]
+    
+    results = df.iloc[top_idx][["somali", "english", "italian", "domain"]].copy()
+    results["similarity_score"] = np.round(scores[top_idx], 4)
+    results["confidence_label"] = [get_confidence_label(s) for s in scores[top_idx]]
+    results.insert(0, "rank", range(1, len(results) + 1))
+    results = results.reset_index(drop=True)
+    
+    return results
+
+if __name__ == "__main__":
+    import sys
+    query = sys.argv[1] if len(sys.argv) > 1 else "dhaqaale"
+    print(f"Searching for: {query}\n")
+    results = search(query)
+    print(results.to_string(index=False))