Update app.py

app.py

@@ -5,6 +5,9 @@ import random
 from sentence_transformers import SentenceTransformer, util
 from datasets import load_dataset
 from spaces import GPU
+import re
+
+
 
 # --- Load model ---
 model = SentenceTransformer("Shuu12121/CodeSearch-ModernBERT-Owl")
@@ -13,23 +16,34 @@ model.eval()
 # --- Load CodeSearchNet dataset (test split only) ---
 dataset = load_dataset("code_x_glue_tc_nl_code_search_adv", trust_remote_code=True, split="test")
 
+def remove_comments_from_code(code: str) -> str:
+    # 複数行コメント（docstring含む）を除去
+    code = re.sub(r'"""[\s\S]*?"""', '', code)
+    code = re.sub(r"'''[\s\S]*?'''", '', code)
+
+    # 単一行コメント（# 以降を除去）
+    code = re.sub(r'#.*', '', code)
+
+    return code
+
+
 # --- Query & Candidate Generator ---
 def get_query_and_candidates(seed: int = 42):
     random.seed(seed)
     idx = random.randint(0, len(dataset) - 1)
     query = dataset[idx]
-    correct_code = query["code"]
+    correct_code = remove_comments_from_code(query["code"])  # 修正
     doc_str = query["docstring"]
 
-    # 正例 + ランダム負例（正例を除く）
     candidate_pool = [example for i, example in enumerate(dataset) if i != idx]
-    negatives = random.sample(candidate_pool, k=99)  # 9件の負例
-    candidates = [correct_code] + [neg["code"] for neg in negatives]
+    negatives = random.sample(candidate_pool, k=99)
+    candidates = [correct_code] + [remove_comments_from_code(neg["code"]) for neg in negatives]  # 修正
     random.shuffle(candidates)
 
     return doc_str, correct_code, candidates
 
 
+
 @GPU
 def code_search_demo(seed: int):
     doc_str, correct_code, candidates = get_query_and_candidates(seed)