Add eval and torch.no_grad (because inference only)

app.py

@@ -2,13 +2,14 @@ from html import escape
 import re
 import streamlit as st
 import pandas as pd, numpy as np
+import torch
 from transformers import CLIPProcessor, CLIPModel
 from st_clickable_images import clickable_images
 
 MODEL_NAMES = [
-#    "base-patch32",
-#    "base-patch16",
-#    "large-patch14",    
+    #    "base-patch32",
+    #    "base-patch16",
+    #    "large-patch14",
     "large-patch14-336"
 ]
 
@@ -20,7 +21,7 @@ def load():
     processors = {}
     embeddings = {}
     for name in MODEL_NAMES:
-        models[name] = CLIPModel.from_pretrained(f"openai/clip-vit-{name}")
+        models[name] = CLIPModel.from_pretrained(f"openai/clip-vit-{name}").eval()
         processors[name] = CLIPProcessor.from_pretrained(f"openai/clip-vit-{name}")
         embeddings[name] = {
             0: np.load(f"embeddings-vit-{name}.npy"),
@@ -39,7 +40,8 @@ source = {0: "\nSource: Unsplash", 1: "\nSource: The Movie Database (TMDB)"}
 
 def compute_text_embeddings(list_of_strings, name):
     inputs = processors[name](text=list_of_strings, return_tensors="pt", padding=True)
-    result = models[name].get_text_features(**inputs).detach().numpy()
+    with torch.no_grad():
+        result = models[name].get_text_features(**inputs).detach().numpy()
     return result / np.linalg.norm(result, axis=1, keepdims=True)
 
 
@@ -158,9 +160,9 @@ def main():
     st.sidebar.markdown(description)
     with st.sidebar.expander("Advanced use"):
         st.markdown(howto)
-    #mode = st.sidebar.selectbox(
+    # mode = st.sidebar.selectbox(
     #    "", ["Results for ViT-L/14@336px", "Comparison of 2 models"], index=0
-    #)
+    # )
 
     _, c, _ = st.columns((1, 3, 1))
     if "query" in st.session_state:
@@ -176,7 +178,7 @@ def main():
         "ViT-L/14@336px (slower)": "large-patch14-336",
     }
 
-    if False:#"Comparison" in mode:
+    if False:  # "Comparison" in mode:
         c1, c2 = st.columns((1, 1))
         selection1 = c1.selectbox("", models_dict.keys(), index=0)
         selection2 = c2.selectbox("", models_dict.keys(), index=2)
@@ -187,7 +189,7 @@ def main():
 
     if len(query) > 0:
         results1 = image_search(query, corpus, name1)
-        if False:#"Comparison" in mode:
+        if False:  # "Comparison" in mode:
             with c1:
                 clicked1 = clickable_images(
                     [result[0] for result in results1],
@@ -225,7 +227,7 @@ def main():
             if change_query:
                 if clicked1 >= 0:
                     st.session_state["query"] = f"[{corpus}:{results1[clicked1][2]}]"
-                #elif clicked2 >= 0:
+                # elif clicked2 >= 0:
                 #    st.session_state["query"] = f"[{corpus}:{results2[clicked2][2]}]"
                 st.experimental_rerun()