Spaces:
Runtime error
Runtime error
File size: 1,542 Bytes
41782f4 84f3fa1 41782f4 84f3fa1 41782f4 e7197be 41782f4 e7197be 41782f4 e7197be 41782f4 2cbb9e5 41782f4 999e0ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import faiss
import joblib
import numpy as np
import pandas as pd
import streamlit as st
from sentence_transformers import SentenceTransformer
@st.experimental_memo
def load_model():
return SentenceTransformer("TamedWicked/MathBERT_hr")
@st.experimental_memo
def load_knowledge_base_df():
return pd.read_parquet("data/knowledge_base.parquet")
@st.experimental_memo
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, 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]]
result.index = list(range(1, len(result)+1))
speeches = result["Speech"].tolist()
links = result["start_link"].tolist()
st.write(show_df_as_markdown(result))
|