Update app.py

app.py

@@ -15,65 +15,84 @@ from PIL import Image
 from txtai.embeddings import Embeddings
 
 
-def images(directory):
+class Application:
     """
-    Generator that loops over each image in a directory.
-
-    Args:
-        directory: directory with images
+    Main application
     """
 
-    for path in glob.glob(directory + "/*jpg") + glob.glob(directory + "/*png"):
-        yield (path, Image.open(path), None)
+    def __init__(self, directory):
+        """
+        Creates a new application.
 
+        Args:
+            directory: directory of images
+        """
 
-@st.cache(allow_output_mutation=True)
-def build(directory):
-    """
-    Builds an image embeddings index.
+        self.embeddings = self.build(directory)
 
-    Args:
-        directory: directory with images
+    def build(self, directory):
+        """
+        Builds an image embeddings index.
 
-    Returns:
-        Embeddings index
-    """
+        Args:
+            directory: directory with images
 
-    embeddings = Embeddings({"method": "sentence-transformers", "path": "clip-ViT-B-32"})
-    embeddings.index(images(directory))
+        Returns:
+            Embeddings index
+        """
 
-    # Update model to support multilingual queries
-    embeddings.config["path"] = "sentence-transformers/clip-ViT-B-32-multilingual-v1"
-    embeddings.model = embeddings.loadVectors()
+        embeddings = Embeddings({"method": "sentence-transformers", "path": "clip-ViT-B-32"})
+        embeddings.index(self.images(directory))
 
-    return embeddings
+        # Update model to support multilingual queries
+        embeddings.config["path"] = "sentence-transformers/clip-ViT-B-32-multilingual-v1"
+        embeddings.model = embeddings.loadVectors()
 
+        return embeddings
 
-def app(directory):
-    """
-    Streamlit application that runs searches against an image embeddings index.
+    def images(self, directory):
+        """
+        Generator that loops over each image in a directory.
 
-    Args:
-        directory: directory with images
-    """
+        Args:
+            directory: directory with images
+        """
 
-    # Build embeddings index
-    embeddings = build(directory)
+        for path in glob.glob(directory + "/*jpg") + glob.glob(directory + "/*png"):
+            yield (path, Image.open(path), None)
 
-    st.title("Image search")
+    def run(self):
+        """
+        Runs a Streamlit application.
+        """
 
-    st.markdown("This application shows how images and text can be embedded into the same space to support similarity search. ")
-    st.markdown(
-        "[sentence-transformers](https://github.com/UKPLab/sentence-transformers/tree/master/examples/applications/image-search) "
-        + "recently added support for the [OpenAI CLIP model](https://github.com/openai/CLIP). This model embeds text and images into "
-        + "the same space, enabling image similarity search. txtai can directly utilize these models."
-    )
+        st.title("Image search")
 
-    query = st.text_input("Search query:")
-    if query:
-        index, _ = embeddings.search(query, 1)[0]
-        st.image(Image.open(index))
+        st.markdown("This application shows how images and text can be embedded into the same space to support similarity search. ")
+        st.markdown(
+            "[sentence-transformers](https://github.com/UKPLab/sentence-transformers/tree/master/examples/applications/image-search) "
+            + "recently added support for the [OpenAI CLIP model](https://github.com/openai/CLIP). This model embeds text and images into "
+            + "the same space, enabling image similarity search. txtai can directly utilize these models."
+        )
+
+        query = st.text_input("Search query:")
+        if query:
+            index, _ = self.embeddings.search(query, 1)[0]
+            st.image(Image.open(index))
+
+@st.cache(allow_output_mutation=True)
+def create(directory):
+    """
+    Creates and caches a Streamlit application.
+
+    Args:
+        directory: directory of images to index
+
+    Returns:
+        Application
+    """
 
+    return Application(directory)
 
 if __name__ == "__main__":
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -86,4 +105,6 @@ if __name__ == "__main__":
         f = tarfile.open(fileobj=response.raw, mode="r|gz")
         f.extractall(path="/tmp")
  
-    app(files)
+    # Create and run application
+    app = create(files)
+    app.run()