Rename ap.py to app.py

ap.py → app.py

@@ -22,17 +22,18 @@ TOKENIZER_PATH = "bpe.model"
 
 if not os.path.exists(MODEL_PATH):
     download_file(
-        "https://huggingface.co/OpenLab-NLP/openlem2/resolve/main/encoder_fit.weights.h5?download=true",
+        "https://huggingface.co/OpenLab-NLP/openlem2-retrieval-qa/resolve/main/encoder_fit.weights.h5?download=true",
         MODEL_PATH
     )
 
 if not os.path.exists(TOKENIZER_PATH):
     download_file(
-        "https://huggingface.co/OpenLab-NLP/openlem2/resolve/main/bpe.model?download=true",
+        "https://huggingface.co/OpenLab-NLP/openlem2-retrieval-qa/resolve/main/bpe.model?download=true",
         TOKENIZER_PATH
     )
 
 MAX_LEN = 384
+TOP_K = 3
 EMBED_DIM = 512
 LATENT_DIM = 512
 BATCH_SIZE = 768           # global batch size (Keras/TPU가 replica-wise로 나눠서 처리)
@@ -190,48 +191,47 @@ encoder = SentenceEncoder(vocab_size=vocab_size)
 encoder(np.zeros((1, MAX_LEN), dtype=np.int32))  # 모델 빌드
 encoder.load_weights(MODEL_PATH)
 
-# ===============================
-# 4️⃣ 벡터화 함수
-# ===============================
-def get_sentence_vector(sentence):
-    tokens = pad_sentence(encode_sentence(sentence))
-    vec = encoder(np.array([tokens])).numpy()[0]
-    return vec / np.linalg.norm(vec)
+def tokenize(texts):
+    token_ids = []
+    for t in texts:
+        ids = sp.encode(t, out_type=int)[:MAX_LEN]
+        if len(ids) < MAX_LEN:
+            ids += [pad_id]*(MAX_LEN-len(ids))
+        token_ids.append(ids)
+    return np.array(token_ids, dtype=np.int32)
 
-# ===============================
-# 5️⃣ 가장 비슷한 문장 찾기
-# ===============================
-def find_most_similar(query, s1, s2, s3):
-    candidates = [s1, s2, s3]
-    candidate_vectors = np.stack([get_sentence_vector(c) for c in candidates]).astype(np.float32)
-    query_vector = get_sentence_vector(query)
-    
-    sims = candidate_vectors @ query_vector  # cosine similarity
-    top_idx = np.argmax(sims)
-    
-    return {
-        "가장 비슷한 문장": candidates[top_idx],
-        "유사도": float(sims[top_idx])
-    }
 
-# ===============================
-# 6️⃣ Gradio UI
-# ===============================
+def search_and_answer(query, docs_text):
+    docs = [d.strip() for d in docs_text.split("\n") if d.strip()]
+    if not docs:
+        return [], "문서를 한 줄씩 입력하세요."
+
+    q_ids = tokenize([query])
+    d_ids = tokenize(docs)
+
+    q_emb = encoder(q_ids, training=False).numpy()
+    d_embs = encoder(d_ids, training=False).numpy()
+
+    scores = np.dot(q_emb, d_embs.T)[0]
+    top_k_idx = scores.argsort()[::-1][:min(TOP_K, len(docs))]
+    top_docs = [(docs[i], float(scores[i])) for i in top_k_idx]
+
+    answer = docs[top_k_idx[0]]
+    return top_docs, answer
+
 with gr.Blocks() as demo:
-    gr.Markdown("## 🔍 문장 유사도 검색기 (쿼리 1개 + 후보 3개)")
-    with gr.Row():
-        query_input = gr.Textbox(label="검색할 문장 (Query)", placeholder="여기에 입력")
+    gr.Markdown("## OpenLEM2 Retrieval-QA 데모 (사용자 문서 입력 가능)")
+
     with gr.Row():
-        s1_input = gr.Textbox(label="검색 후보 1")
-        s2_input = gr.Textbox(label="검색 후보 2")
-        s3_input = gr.Textbox(label="검색 후보 3")
-    output = gr.JSON(label="결과")
+        query_input = gr.Textbox(label="질문/쿼리", placeholder="예: 서울 날씨 어때?")
     
-    search_btn = gr.Button("가장 비슷한 문장 찾기")
-    search_btn.click(
-        fn=find_most_similar,
-        inputs=[query_input, s1_input, s2_input, s3_input],
-        outputs=output
-    )
+    docs_input = gr.Textbox(label="문서 리스트 (한 줄씩)", placeholder="문서를 한 줄씩 입력하세요.", lines=10)
+
+    with gr.Row():
+        top_docs_out = gr.Dataframe(headers=["Document", "Score"], max_rows=TOP_K)
+        answer_out = gr.Textbox(label="답변")
+
+    run_btn = gr.Button("검색/QA 실행")
+    run_btn.click(fn=search_and_answer, inputs=[query_input, docs_input], outputs=[top_docs_out, answer_out])
 
 demo.launch()