Add app

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.9 as core
+
+COPY ./core/requirements.txt ./requirements.txt
+RUN pip install -f requirements.txt
+
+WORKDIR /app
+
+COPY ./core .
+RUN python ./initialization.py
+
+FROM python:3.9
+
+COPY ./api/requirements.txt ./requirements.txt
+RUN pip install -f requirements.txt
+
+WORKDIR /app
+COPY ./api .
+
+COPY --from=core ./engine.pickle ./engine.pickle
+
+EXPOSE 9999
+ENTRYPOINT ["python", "service_manager.py"]

api/__init__.py

@@ -0,0 +1 @@
+# This file is intentionally left empty

api/requirements.txt

@@ -0,0 +1,3 @@
+fastapi==0.111.1
+pydantic==2.8.2
+uvicorn==0.30.3

api/service_manager.py

@@ -0,0 +1,13 @@
+from web_server import app
+import uvicorn
+
+
+def run():
+    """
+    TODO
+    """
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+
+run()
+

api/web_server.py

@@ -0,0 +1,26 @@
+import dill
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from core.search_engine import PromptSearchEngine
+
+
+class Query(BaseModel):
+    prompt: str
+    n: int = 5
+
+
+app = FastAPI()
+
+with open('../core/engine.pickle', 'rb') as file:
+    serialized_engine = file.read()
+
+prompt_search_engine = dill.loads(serialized_engine)
+
+
+@app.post("/search/")
+async def search(query: Query):
+    try:
+        results = prompt_search_engine.most_similar(query.prompt, query.n)
+        return {"results": results}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

core/__init__.py

@@ -0,0 +1 @@
+# This file is intentionally left empty

core/data/__init__.py

@@ -0,0 +1 @@
+# This file is intentionally left empty

core/data/dataset.py

@@ -0,0 +1,44 @@
+from datasets import load_dataset
+from core.search_engine import PromptSearchEngine
+
+
+class PromptDataset:
+    """
+    TODO
+    """
+
+    def __init__(self, dataset_name: str):
+        """
+        TODO
+        """
+
+        self.dataset_name = dataset_name
+        self.dataset = None
+
+    def load(self):
+        """
+        TODO
+        """
+
+        self.dataset = load_dataset(self.dataset_name)
+
+        return self.dataset
+
+    def get_prompts(self):
+        """
+        TODO
+        """
+
+        if self.dataset is None:
+            raise ValueError("Dataset not loaded. Call the load() method first.")
+
+        return [item['Prompt'] for item in self.dataset['test']]
+
+
+# if __name__ == "__main__":
+#     dataset = PromptDataset("Gustavosta/Stable-Diffusion-Prompts")
+#     dataset.load()
+#     prompts = dataset.get_prompts()
+#     engine = PromptSearchEngine(prompts)
+#     result = engine.most_similar("dark")
+#     print(result)

core/initialization.py

@@ -0,0 +1,21 @@
+import dill
+from data.dataset import PromptDataset
+from core.search_engine import PromptSearchEngine
+
+
+def run():
+    """
+    TODO
+    """
+
+    prompt_dataset = PromptDataset("Gustavosta/Stable-Diffusion-Prompts")
+    prompt_dataset.load()
+    prompts = prompt_dataset.get_prompts()
+    engine = PromptSearchEngine(prompts)
+
+    serialized_engine = dill.dumps(engine)
+
+    with open("engine.pickle", "wb") as file:
+        file.write(serialized_engine)
+
+run()

core/requirements.txt

@@ -0,0 +1,3 @@
+datasets==2.20.0
+faiss_cpu==1.8.0.post1
+sentence_transformers==3.0.1

core/search_engine.py

@@ -0,0 +1,36 @@
+from typing import List, Sequence, Tuple
+import numpy as np
+import faiss
+from core.vectorizer import Vectorizer
+
+
+class PromptSearchEngine(object):
+    """
+    TODO
+    """
+
+    def __init__(self, prompts: Sequence[str]) -> None:
+        """
+        TODO
+        """
+
+        self.vectorizer = Vectorizer()
+        self.corpus_vectors = self.vectorizer.transform(prompts)
+        self.corpus = prompts
+
+        self.corpus_vectors = self.corpus_vectors / np.linalg.norm(self.corpus_vectors, axis=1, keepdims=True)
+
+        d = self.corpus_vectors.shape[1]
+        self.index = faiss.IndexFlatIP(d)
+        self.index.add(self.corpus_vectors.astype('float32'))
+
+    def most_similar(self, query: str, n: int = 5) -> List[Tuple[float, str]]:
+        """
+        TODO
+        """
+
+        query_vector = self.vectorizer.transform([query]).astype('float32')
+        query_vector = query_vector / np.linalg.norm(query_vector)
+        distances, indices = self.index.search(query_vector, n)
+
+        return [(distances[0][i], self.corpus[indices[0][i]]) for i in range(n)]

core/vectorizer.py

@@ -0,0 +1,34 @@
+from typing import Sequence
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+
+class Vectorizer(object):
+    """
+    TODO
+    """
+
+    def __init__(self, model_name: str = 'all-MiniLM-L6-v2') -> None:
+        """
+        Initialize the vectorizer with a pre-trained embedding model.
+        """
+
+        self.model = SentenceTransformer(model_name)
+
+    def transform(self, prompts: Sequence[str]) -> np.ndarray:
+        """
+        Transform texts into numerical vectors using the specified model.
+        """
+
+        return self.model.encode(list(prompts))
+
+    @staticmethod
+    def cosine_similarity(query_vector: np.ndarray, corpus_vectors: np.ndarray) -> np.ndarray:
+        """
+        Calculate cosine similarity between prompt vectors.
+        """
+        
+        query_norm = query_vector / np.linalg.norm(query_vector)
+        corpus_norms = corpus_vectors / np.linalg.norm(corpus_vectors, axis=1, keepdims=True)
+
+        return np.dot(corpus_norms, query_norm.T).flatten()