Added knowledge bases and simple query with top 5 results as a table.

app.py

@@ -1,4 +1,51 @@
+import faiss
+import joblib
+import numpy as np
+import pandas as pd
 import streamlit as st
+from sentence_transformers import SentenceTransformer
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+# st.set_page_config(layout="wide")
+
+
+@st.cache_resource
+def load_model():
+    return SentenceTransformer("TamedWicked/MathBERT_hr")
+
+
+@st.cache_resource
+def load_knowledge_base_df():
+    return pd.read_parquet("data/knowledge_base.parquet")
+
+
+@st.cache_resource
+def load_knowledge_base_index():
+    embeddings = joblib.load("data/knowledge_base_embeddings.pkl")
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
+    return index
+
+
+def vector_search(query: list[str], model: SentenceTransformer, index, num_results=10):
+    vector = model.encode(list(query), show_progress_bar=False, convert_to_numpy=True)
+    D, I = index.search(np.array(vector).astype("float32"), k=num_results)
+    return D, I
+
+
+def show_df_as_html(df: pd.DataFrame):
+    return df.to_html()
+
+
+def show_df_as_markdown(df: pd.DataFrame):
+    return df.to_markdown()
+
+
+model: SentenceTransformer = load_model()
+df: pd.DataFrame = load_knowledge_base_df()
+knowledge_index: np.array = load_knowledge_base_index()
+
+query = st.text_input("Your math query:", value="Jesu li strukture koje su elementarno ekvivalentne izomorfne?")
+if query:
+    D, I = vector_search([query], model, knowledge_index, num_results=5)
+    result = df[["Speech", "start_link"]].iloc[I[0]]
+    st.write(show_df_as_markdown(result), unsafe_allow_html=True)

data/knowledge_base.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b60bdb181b8841e6d0d2e56f645c7070bf1017d056be4950f0facca9936be2a
+size 999177

data/knowledge_base_embeddings.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bc7c3a5a91b52295c08bd22289a11d609f9a273e760c381041859bebe0c0223
+size 58875105