Spaces:
Runtime error
Runtime error
Commit
·
cf4befe
1
Parent(s):
609cb5d
Update app.py
Browse files
app.py
CHANGED
@@ -32,7 +32,7 @@ st.cache(show_spinner=False)
|
|
32 |
def image_search(query, top_k=10):
|
33 |
with torch.no_grad():
|
34 |
text_embedding = text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
|
35 |
-
|
36 |
return [links[i] for i in indices[:top_k]]
|
37 |
|
38 |
|
|
|
32 |
def image_search(query, top_k=10):
|
33 |
with torch.no_grad():
|
34 |
text_embedding = text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
|
35 |
+
_, indices = torch.cosine_similarity(image_embeddings, text_embedding).sort(descending=True)
|
36 |
return [links[i] for i in indices[:top_k]]
|
37 |
|
38 |
|