Add embedding generator

.gitignore

@@ -109,6 +109,7 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+lex-semantic-search/
 
 # Spyder project settings
 .spyderproject
@@ -127,3 +128,10 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# IDE
+.idea
+*.iml
+
+# Custom files
+data/

README.md

@@ -1,2 +1,14 @@
 # lex-semantic-search
 Semantic search for Lex Fridman podcast
+
+
+## Dataset
+
+## Usage
+
+```bash
+python -m venv lex-semantic-search
+source lex-semantic-search/bin/activate 
+pip install -r requirements_cpu.txt # for CPU
+pip install -r requirements_gpu.txt # for GPU
+```

embeddings/encoder.py

@@ -0,0 +1,9 @@
+from abc import ABC, abstractmethod
+from typing import List
+import numpy as np
+
+
+class EmbeddingEncoder(ABC):
+    @abstractmethod
+    def generate_embeddings(self, texts: List[str]) -> List[np.ndarray]:
+        pass

embeddings/huggingface.py

@@ -0,0 +1,30 @@
+import torch
+import numpy as np
+
+from typing import List
+from transformers import AutoTokenizer, AutoModel
+from embeddings.encoder import EmbeddingEncoder
+
+
+def cls_pooling(model_output):
+    return model_output.last_hidden_state[:, 0]
+
+
+class HuggingFaceEncoder(EmbeddingEncoder):
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+        self.model = AutoModel.from_pretrained(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def generate_embeddings(self, sentences: List[str]) -> List[np.ndarray]:
+        # Tokenize sentences
+        encoded_input = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
+
+        # Compute token embeddings
+        with torch.no_grad():
+            model_output = self.model(**encoded_input, return_dict=True)
+
+        # Perform pooling
+        embeddings = cls_pooling(model_output)
+
+        return embeddings