first commit

app.py

@@ -0,0 +1,84 @@
+import json
+import gradio as gr
+import torch
+from sklearn.metrics.pairwise import cosine_similarity
+from transformers import AutoTokenizer, AutoModel, MarianMTModel, MarianTokenizer
+
+# === 模型加载 ===
+print("Loading models...")
+embed_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-large-en")
+embed_model = AutoModel.from_pretrained("BAAI/bge-large-en")
+
+trans_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
+trans_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
+print("Models loaded.")
+
+# === 载入字典文件 ===
+with open("k_definition_cleaned.json", encoding="utf-8") as f:
+    modern_dict = json.load(f)
+
+with open("oc_definition_cleaned.json", encoding="utf-8") as f:
+    ancient_dict = json.load(f)
+
+# === 编码函数 ===
+def encode(texts):
+    inputs = embed_tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        outputs = embed_model(**inputs)
+        embeddings = outputs.last_hidden_state[:, 0]
+        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+    return embeddings
+
+# === 翻译函数 ===
+def translate_to_english(text):
+    if all(ord(c) < 128 for c in text):  # already English
+        return text
+    inputs = trans_tokenizer(text, return_tensors="pt", padding=True)
+    translated = trans_model.generate(**inputs)
+    return trans_tokenizer.decode(translated[0], skip_special_tokens=True)
+
+# === 核心匹配函数 ===
+def find_similar_hanzi(idea_text, top_k=10):
+    idea_en = translate_to_english(idea_text)
+    idea_vec = encode([idea_en])[0].unsqueeze(0)
+
+    def search(dictionary):
+        results = []
+        for hanzi, defs in dictionary.items():
+            def_vecs = encode(defs)
+            scores = cosine_similarity(def_vecs, idea_vec).flatten()
+            max_idx = scores.argmax()
+            results.append((hanzi, defs[max_idx], float(scores[max_idx])))
+        return sorted(results, key=lambda x: x[2], reverse=True)[:top_k]
+
+    modern = search(modern_dict)
+    ancient = search(ancient_dict)
+
+    return modern, ancient
+
+# === 用于展示表格的处理函数 ===
+def gradio_interface(query):
+    modern, ancient = find_similar_hanzi(query, top_k=50)
+    return {
+        "modern_results": [[h, d, round(s, 4)] for h, d, s in modern],
+        "ancient_results": [[h, d, round(s, 4)] for h, d, s in ancient]
+    }
+
+# === Gradio 页面设置 ===
+with gr.Blocks() as demo:
+    gr.Markdown("# Hanzi Imagery Search")
+    with gr.Row():
+        inp = gr.Textbox(label="输入意象短语（中/英文）", placeholder="如：warrior, warmth, 月亮等")
+        btn = gr.Button("搜索")
+
+    modern_output = gr.Dataframe(headers=["汉字", "释义", "相似度"], label="现代释义匹配", interactive=False)
+    ancient_output = gr.Dataframe(headers=["汉字", "释义", "相似度"], label="古代释义匹配", interactive=False)
+    json_output = gr.JSON(label="JSON 返回结构")
+
+    def full_response(query):
+        res = gradio_interface(query)
+        return res["modern_results"], res["ancient_results"], res
+
+    btn.click(fn=full_response, inputs=[inp], outputs=[modern_output, ancient_output, json_output])
+
+demo.launch(share=True)