Upload 14 files

Dockerfile

@@ -0,0 +1,11 @@
+# Dockerfile
+FROM python:3.9-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "run:app", "--host", "0.0.0.0", "--port", "8000"]

app/search_engine.py

@@ -0,0 +1,54 @@
+import numpy as np
+from typing import List, Tuple
+
+from .similarity import cosine_similarity
+from .vectorizer import Vectorizer
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class PromptSearchEngine:
+    def __init__(self):
+        self.vectorizer = Vectorizer(init_pinecone=False)
+        self.vectorizer._data_loaded = True
+        self.prompts = self.vectorizer.prompts
+        self.corpus_vectors = self.vectorizer.transform(self.prompts)
+        self.index_name = self.vectorizer.pinecone_index_name
+
+    def most_similar(self, query: str, n: int = 5, use_pinecone=True) -> List[Tuple[float, str]]:
+        logger.info(f"Encoding query: {query}")
+        query_vector = self.vectorizer.transform([query])[0]
+        logger.info(f"Encoded query vector: {query_vector}")
+        if use_pinecone:
+            logger.info(f"I'm doing pinecone vector search because the use_pinecone is: {use_pinecone}")
+            try:
+                # Convert numpy array to list of native Python floats
+                query_vector_list = query_vector.tolist()
+                search_result = self.vectorizer.index.query(
+                    vector=query_vector_list,
+                    top_k=n,
+                    include_metadata=True
+                )
+                logger.info(f"Search result: {search_result}")
+
+                # Retrieve and format the results
+                results = [(match['score'], match['metadata']['text']) for match in search_result['matches'] if
+                           'text' in match['metadata']]
+            except Exception as e:
+                logger.error(f"Pinecone query failed: {e}")
+                logger.info("Falling back to cosine similarity search.")
+
+                # Fallback to cosine similarity search
+                similarities = cosine_similarity(query_vector, self.corpus_vectors)
+                top_n_indices = np.argsort(similarities)[-n:][::-1]
+                results = [(float(similarities[i]), self.prompts[i]) for i in top_n_indices]
+        else:
+            logger.info(f"I'm cosine similarity search because the use_pinecone is: {use_pinecone}")
+            logger.info("Using cosine similarity for search")
+            similarities = cosine_similarity(query_vector, self.corpus_vectors)
+            top_n_indices = np.argsort(similarities)[-n:][::-1]
+            results = [(float(similarities[i]), self.prompts[i]) for i in top_n_indices]
+        return results

app/similarity.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+
+def cosine_similarity(
+    query_vector: np.ndarray,
+    corpus_vectors: np.ndarray
+) -> np.ndarray:
+    """
+    Calculate cosine similarity between a query vector and a corpus of vectors.
+
+    Args:
+        query_vector: Vectorized prompt query of shape (D,).
+        corpus_vectors: Vectorized prompt corpus of shape (N, D).
+
+    Returns:
+        np.ndarray: The vector of shape (N,) with values in range [-1, 1] where 1
+        is max similarity i.e., two vectors are the same.
+    """
+    dot_product = np.dot(corpus_vectors, query_vector)
+    norm_query = np.linalg.norm(query_vector)
+    norm_corpus = np.linalg.norm(corpus_vectors, axis=1)
+    return dot_product / (norm_query * norm_corpus)

app/vectorizer.py

@@ -0,0 +1,98 @@
+import os
+import logging
+
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from datasets import load_dataset
+from pinecone import Pinecone, ServerlessSpec
+
+# Disable parallelism for tokenizers
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class Vectorizer:
+    def __init__(self, model_name='all-mpnet-base-v2', batch_size=64, init_pinecone=True):
+        logger.info(f"Initializing Vectorizer with model {model_name} and batch size {batch_size}")
+        self.model = SentenceTransformer(model_name)
+        self.prompts = []
+        self.batch_size = batch_size
+        self.pinecone_index_name = "prompts-index"
+        self._init_pinecone = init_pinecone
+        self._setup_pinecone()
+        self._load_prompts()
+
+    def _setup_pinecone(self):
+        logger.info("Setting up Pinecone")
+        # Initialize Pinecone
+        pinecone = Pinecone(api_key='b514eb66-8626-4697-8a1c-4c411c06c090')
+        # Check if the Pinecone index exists, if not create it
+        existing_indexes = pinecone.list_indexes()
+
+        logger.info(f"self.init_pineconeself.init_pineconeself"
+                    f".init_pineconeself.init_pineconeself.init_pinecone: {self._init_pinecone}")
+        if self.pinecone_index_name not in existing_indexes:
+            logger.info(f"Creating Pinecone index: {self.pinecone_index_name}")
+            if self._init_pinecone:
+                pinecone.create_index(
+                    name=self.pinecone_index_name,
+                    dimension=768,
+                    metric='cosine',
+                    spec=ServerlessSpec(
+                        cloud="aws",
+                        region="us-east-1"
+                    )
+                )
+        else:
+            logger.info(f"Pinecone index {self.pinecone_index_name} already exists")
+
+        self.index = pinecone.Index(self.pinecone_index_name)
+
+    def _load_prompts(self):
+        logger.info("Loading prompts from Pinecone")
+        self.prompts = []
+        # Fetch vectors from the Pinecone index
+        index_stats = self.index.describe_index_stats()
+        logger.info(f"Index stats: {index_stats}")
+
+        namespaces = index_stats['namespaces']
+        for namespace, stats in namespaces.items():
+            vector_count = stats['vector_count']
+            ids = [str(i) for i in range(vector_count)]
+            for i in range(0, vector_count, self.batch_size):
+                batch_ids = ids[i:i + self.batch_size]
+                response = self.index.fetch(ids=batch_ids)
+                for vector in response.vectors.values():
+                    metadata = vector.get('metadata')
+                    if metadata and 'text' in metadata:
+                        self.prompts.append(metadata['text'])
+        logger.info(f"Loaded {len(self.prompts)} prompts from Pinecone")
+
+    def _store_prompts(self, dataset):
+        logger.info("Storing prompts in Pinecone")
+        for i in range(0, len(dataset), self.batch_size):
+            batch = dataset[i:i + self.batch_size]
+            vectors = self.model.encode(batch)
+            # Prepare data for Pinecone
+            pinecone_data = [{'id': str(i + j), 'values': vector.tolist(), 'metadata': {'text': batch[j]}} for j, vector
+                             in enumerate(vectors)]
+            self.index.upsert(vectors=pinecone_data)
+            logger.info(f"Upserted batch {i // self.batch_size + 1}/{len(dataset) // self.batch_size + 1} to Pinecone")
+
+    def transform(self, prompts):
+        return np.array(self.model.encode(prompts))
+
+    def store_from_dataset(self, store_data=False):
+        if store_data:
+            logger.info("Loading dataset")
+            dataset = load_dataset('fantasyfish/laion-art', split='train')
+            logger.info(f"Loaded {len(dataset)} items from dataset")
+            logger.info("Please wait for storing. This may take up to five minutes. ")
+            self._store_prompts([item['text'] for item in dataset])
+            logger.info("Items from dataset are stored.")
+            # Ensure prompts are loaded after storing
+            self._load_prompts()
+            logger.info("Items from dataset are loaded.")

load_data.py

@@ -0,0 +1,5 @@
+from app.vectorizer import Vectorizer
+
+if __name__ == "__main__":
+    vectorizer = Vectorizer()
+    vectorizer.store_from_dataset(store_data=True)  # Run this once to load the dataset into Pinecone

requirements.txt

@@ -0,0 +1,9 @@
+flask
+requests
+streamlit
+transformers
+numpy
+sentence-transformers
+datasets
+pinecone
+unittest

run.py

@@ -0,0 +1,33 @@
+import os
+import logging
+from flask import Flask, request, jsonify
+from app.search_engine import PromptSearchEngine
+
+app = Flask(__name__)
+
+# Disable parallelism for tokenizers
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+search_engine = PromptSearchEngine()
+
+@app.route('/search', methods=['POST'])
+def search():
+    data = request.get_json()
+    query = data.get('query')
+    n = data.get('n', 5)
+    use_pinecone = data.get('use_pinecone', True)
+
+    logger.info(f"Received query: {query} with n: {n} and use_pinecone: {use_pinecone}")
+    results = search_engine.most_similar(query, n, use_pinecone)
+    formatted_results = [{'score': score, 'prompt': prompt} for score, prompt in results]
+    logger.info(f"Returning results: {formatted_results}")
+    return jsonify(formatted_results)
+
+
+if __name__ == '__main__':
+    logger.info("Starting Flask server")
+    app.run(debug=True)

run_tests.py

@@ -0,0 +1,13 @@
+import unittest
+
+
+def run_all_tests():
+    test_loader = unittest.TestLoader()
+    test_suite = test_loader.discover('tests', pattern='test_*.py')
+
+    test_runner = unittest.TextTestRunner(verbosity=2)
+    test_runner.run(test_suite)
+
+
+if __name__ == '__main__':
+    run_all_tests()

tests/test_search_engine.py

@@ -0,0 +1,39 @@
+import unittest
+from unittest.mock import patch
+import numpy as np
+from app.search_engine import PromptSearchEngine
+
+
+class TestPromptSearchEngine(unittest.TestCase):
+
+    @patch('app.vectorizer.Vectorizer')
+    def setUp(self, mock_vectorizer):
+        self.mock_vectorizer = mock_vectorizer.return_value
+        self.mock_vectorizer.transform.return_value = np.random.rand(10, 768)
+        self.mock_vectorizer.prompts = ['prompt'] * 10
+        self.search_engine = PromptSearchEngine()
+
+    def test_most_similar_with_cosine_similarity(self):
+        self.mock_vectorizer.index.query.side_effect = Exception('Pinecone error')
+        results = self.search_engine.most_similar('query', use_pinecone=False)
+        self.assertEqual(len(results), 5)
+        self.assertIsInstance(results[0][0], float)
+        self.assertIsInstance(results[0][1], str)
+
+    def test_most_similar_with_pinecone(self):
+        mock_search_result = {
+            'matches': [
+                {'score': np.float32(0.9), 'metadata': {'text': 'prompt1'}},
+                {'score': np.float32(0.8), 'metadata': {'text': 'prompt2'}}
+            ]
+        }
+        self.mock_vectorizer.index.query.return_value = mock_search_result
+
+        results = self.search_engine.most_similar('query', use_pinecone=True)
+        self.assertEqual(len(results), 5)
+        self.assertIsInstance(results[0][0], float)
+        self.assertIsInstance(results[0][1], str)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_similarity.py

@@ -0,0 +1,23 @@
+import unittest
+import numpy as np
+from app.similarity import cosine_similarity
+
+
+class TestSimilarity(unittest.TestCase):
+
+    def test_cosine_similarity(self):
+        query_vector = np.array([1, 2, 3])
+        corpus_vectors = np.array([
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9]
+        ])
+
+        expected_result = np.array([1.0, 0.9746318461970762, 0.9594119455666703])
+        result = cosine_similarity(query_vector, corpus_vectors)
+
+        np.testing.assert_almost_equal(result, expected_result, decimal=6)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_vectorizer.py

@@ -0,0 +1,36 @@
+import unittest
+from unittest.mock import patch, MagicMock
+import numpy as np
+
+from app.vectorizer import Vectorizer
+
+
+class TestVectorizer(unittest.TestCase):
+
+    @patch('app.vectorizer.Pinecone')
+    @patch('app.vectorizer.SentenceTransformer')
+    def test_vectorizer_initialization(self, mock_sentence_transformer, mock_pinecone):
+        mock_sentence_transformer.return_value.encode.return_value = np.random.rand(1, 768)
+        vectorizer = Vectorizer(init_pinecone=False)
+        self.assertEqual(vectorizer.batch_size, 64)
+        self.assertEqual(vectorizer.pinecone_index_name, "prompts-index")
+
+    @patch('app.vectorizer.load_dataset')
+    @patch('app.vectorizer.Pinecone')
+    def test_store_from_dataset(self, mock_pinecone, mock_load_dataset):
+        mock_pinecone_instance = MagicMock()
+        mock_pinecone.return_value = mock_pinecone_instance
+        mock_load_dataset.return_value = [{'text': 'sample text'}]
+
+        vectorizer = Vectorizer(init_pinecone=False)
+        vectorizer.store_from_dataset(store_data=True)
+
+        mock_load_dataset.assert_called_once_with('fantasyfish/laion-art', split='train')
+        mock_pinecone_instance.Index.return_value.upsert.assert_called()
+
+    def test_transform(self):
+        with patch('app.vectorizer.SentenceTransformer') as mock_sentence_transformer:
+            mock_sentence_transformer.return_value.encode.return_value = np.random.rand(1, 768)
+            vectorizer = Vectorizer(init_pinecone=False)
+            vectors = vectorizer.transform(['sample prompt'])
+            self.assertEqual(vectors.shape, (1, 768))

ui/app.py

@@ -0,0 +1,32 @@
+import json
+
+import streamlit as st
+import requests
+
+st.title("Prompt Search Engine")
+
+query = st.text_input("Enter your query:")
+use_pinecone = st.radio(
+    "Choose search method:",
+    ('Pinecone Vector Search', 'Cosine Similarity')
+)
+n = st.number_input("Number of results:", min_value=1, max_value=20, value=5)
+
+if st.button("Search"):
+    search_method = use_pinecone == 'Pinecone Vector Search'
+    response = requests.post("http://localhost:5000/search", json={"query": query, "n": n, "use_pinecone": search_method})
+
+    # Log the response for debugging
+    st.write("Response Status Code:", response.status_code)
+
+    try:
+        results = response.json()
+        # for score, prompt in results:
+        #     st.write(f"{score:.2f} - {prompt}")
+        for result in results:
+            score = float(result['score'])
+            prompt = result['prompt']
+            st.write(f"{score:.2f} - {prompt}")
+    except json.JSONDecodeError as e:
+        st.error(f"Failed to decode JSON response: {e}")
+        st.write(response.content)