faiss added

app.py

@@ -6,6 +6,7 @@ import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModel
 from joblib import load
+import faiss
 
 tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
 model = AutoModel.from_pretrained("cointegrated/rubert-tiny2")
@@ -22,29 +23,33 @@ def embed_bert_cls(text, model, tokenizer):
     return embeddings[0].cpu().numpy()
 
 embeded_list = load('embeded_list.joblib')
+index = faiss.IndexFlatL2(embeded_list.shape[1])
+index.add(embeded_list.astype('float32'))
+
 text = st.text_input('Введите текст')
 count_visible = st.number_input("Введите количество отображаемых элементов", 1, 10, 5, step=1)
 if st.button("Найти", type="primary"):
     st.write('Количество фильмов в выборке 4950')
     if text and count_visible:
         embeded_text = embed_bert_cls(text, model, tokenizer).reshape(1,-1)
-        cossim = pairwise_distances(embeded_text, embeded_list)[0]
+        D, I = index.search(embeded_text, index.ntotal)
+        # cossim = pairwise_distances(embeded_text, embeded_list)[0]
         for i in range(count_visible):
             col1, col2 = st.columns(2)
             with col1:
-                st.header(films.iloc[cossim.argsort()].iloc[i][2])
-                st.write(films.iloc[cossim.argsort()].iloc[i][3].replace('\xa0', ' '))
-                st.write(f'Мера схожести евклидова расстояния {sorted(cossim)[i]:4f}')
+                st.header(films.iloc[I[0]].iloc[i][2])
+                st.write(films.iloc[I[0]].iloc[i][3].replace('\xa0', ' '))
+                st.write(f'Мера схожести евклидова расстояния {D[0][i]:4f}')
             with col2:
                 try:
-                    st.image(films.iloc[cossim.argsort()].iloc[i][1])
+                    st.image(films.iloc[I[0]].iloc[i][1])
                 except:
                     st.write('Нет картинки')
         st.header('Самый не подходящий запрос')
         col3, col4 = st.columns(2)
         with col3:
-            st.header(films.iloc[cossim.argsort()].iloc[-1][2])
-            st.write(films.iloc[cossim.argsort()].iloc[-1][3].replace('\xa0', ' '))
-            st.write(f'Мера схожести евклидова расстояния {sorted(cossim)[i]:.4f}')
+            st.header(films.iloc[I[0]].iloc[-1][2])
+            st.write(films.iloc[I[0]].iloc[-1][3].replace('\xa0', ' '))
+            st.write(f'Мера схожести евклидова расстояния {D[0][i]:.4f}')
         with col4:
-            st.image(films.iloc[cossim.argsort()].iloc[-1][1])
+            st.image(films.iloc[I[0]].iloc[-1][1])

requirements.txt

@@ -10,6 +10,7 @@ charset-normalizer==3.2.0
 click==8.1.7
 datasets==2.14.5
 dill==0.3.7
+faiss-cpu==1.7.4
 filelock==3.12.4
 frozenlist==1.4.0
 fsspec==2023.6.0