WIP

app.py

@@ -0,0 +1,77 @@
+import os
+
+import streamlit as st
+from yasem import SpladeEmbedder
+
+if os.getenv("SPACE_ID"):
+    USE_HF_SPACE = True
+    os.environ["HF_HOME"] = "/data/.huggingface"
+    os.environ["HF_DATASETS_CACHE"] = "/data/.huggingface"
+else:
+    USE_HF_SPACE = False
+
+MODEL_NAME = os.environ.get("MODEL_NAME", "hotchpotch/wip-tmp-model-base-v1-pre-2")
+
+
+@st.cache_resource
+def get_embedder(model_name: str = MODEL_NAME) -> SpladeEmbedder:
+    revision = None
+    if MODEL_NAME == "hotchpotch/wip-tmp-model-base-v1-pre-2":
+        revision = "a7db67721ea22165faefba6f0b7ee726b0f3a78f"
+    embedder = SpladeEmbedder(
+        model_name,
+        revision=revision,
+    )
+    return embedder
+
+
+def get_token_values_sorted(input_text: str) -> list[tuple[float, str]]:
+    embedder = get_embedder()
+    embeddings = embedder.encode([input_text])
+    token_values = embedder.get_token_values(embeddings[0])
+    sorted_tokens = sorted(token_values.items(), key=lambda item: item[1], reverse=True)  # type: ignore
+    return [(value, key) for key, value in sorted_tokens]
+
+
+def main():
+    st.set_page_config(
+        page_title="SPLADE 日本語 demo",
+        layout="centered",
+        initial_sidebar_state="auto",
+    )
+
+    st.title("SPLADE 日本語 demo")
+
+    st.markdown("""
+    ### 入力
+    以下のテキストエリアに解析したいテキストを入力してください。
+    """)
+
+    input_text = st.text_area("テキスト入力", height=200)
+
+    if st.button("解析開始"):
+        if input_text.strip():
+            with st.spinner("解析中..."):
+                sorted_tokens = get_token_values_sorted(input_text)
+
+            st.success("解析が完了しました。")
+
+            st.markdown("### 結果")
+            if sorted_tokens:
+                formatted_data = [
+                    {"頻度": freq, "単語": word} for freq, word in sorted_tokens
+                ]
+                st.table(formatted_data)
+            else:
+                st.warning("入力テキストから有効な単語が見つかりませんでした。")
+        else:
+            st.warning("テキストを入力してください。")
+    else:
+        get_embedder()
+
+    st.markdown("---")
+    st.markdown("© 2024 SPLADE 日本語 demo")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1 @@
+yasem