embedding-flow 0.1.0__py3-none-any.whl

contracts/contracts.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+class transform_data(ABC):
+    @abstractmethod
+    def transform_data(self, url: str) -> Optional[str]:
+        """Transforma datos y retorna la ruta del archivo procesado, o None si falla"""
+        pass
+
+class load_data(ABC):
+    @abstractmethod
+    def load_data(self, url: str) -> bool:
+        """Carga datos y retorna True si fue exitoso, False si falló"""
+        pass

embedding_flow-0.1.0.dist-info/METADATA

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: embedding-flow
+Version: 0.1.0
+Summary: Pipeline to transform text chunks into embeddings and load to Qdrant
+Author: facuvega
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: pyarrow>=12.0.0
+Requires-Dist: sentence-transformers>=2.2.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: qdrant-client>=1.7.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Dynamic: license-file

embedding_flow-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+contracts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+contracts/contracts.py,sha256=im3wzTHMxZFVIVyjVVN1u4eZ5J_HxYGrcu3oQ8BXoOk,441
+embedding_flow-0.1.0.dist-info/licenses/LICENSE,sha256=Vi3cItkblr6fZwGbNlp_HnBaMFwXSWYPkrVQLXX3LCs,1057
+load/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+load/load.py,sha256=Y0JbC1x7SfG_EzsSDacDK8fuIcv7K1guQ8AD0birheE,3706
+transform/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+transform/transform.py,sha256=GPeea6D2EbdNwmxdCDtDGgTGIPVLs_Z4VePYQJfzROc,1585
+embedding_flow-0.1.0.dist-info/METADATA,sha256=uFM06vDN-pagU7V_GoZ1wNU6ZPsX8MzdlxNMjV2JoaA,802
+embedding_flow-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+embedding_flow-0.1.0.dist-info/top_level.txt,sha256=VfSCJdxLtTjvyUSe1z1kQgIf0j2zFt3--FivJqWMxxA,25
+embedding_flow-0.1.0.dist-info/RECORD,,

embedding_flow-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

embedding_flow-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2025
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

embedding_flow-0.1.0.dist-info/top_level.txt

@@ -0,0 +1,3 @@
+contracts
+load
+transform

load/load.py

@@ -0,0 +1,101 @@
+from contracts.contracts import load_data
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams, PointStruct
+import pandas as pd
+import logging
+import os
+from typing import List
+import uuid
+
+# Configurar logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+class load_embedding(load_data):
+    def __init__(self):
+        """Inicializa el cliente de Qdrant"""
+        self.qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
+        self.qdrant_api_key = os.getenv("QDRANT_API_KEY", None)
+        self.collection_name = os.getenv("QDRANT_COLLECTION", "embeddings_collection")
+        self.vector_size = int(os.getenv("VECTOR_SIZE", "768"))
+        
+        # Inicializar cliente
+        self.client = QdrantClient(
+            url=self.qdrant_url,
+            api_key=self.qdrant_api_key
+        )
+        
+        # Crear colección si no existe
+        self._ensure_collection_exists()
+    
+    def _ensure_collection_exists(self):
+        """Crea la colección en Qdrant si no existe"""
+        try:
+            collections = self.client.get_collections().collections
+            collection_exists = any(col.name == self.collection_name for col in collections)
+            
+            if not collection_exists:
+                self.client.create_collection(
+                    collection_name=self.collection_name,
+                    vectors_config=VectorParams(
+                        size=self.vector_size,
+                        distance=Distance.COSINE
+                    )
+                )
+                logger.info(f"✅ Colección '{self.collection_name}' creada en Qdrant")
+            else:
+                logger.info(f"ℹ️ Colección '{self.collection_name}' ya existe")
+        except Exception as e:
+            logger.error(f"❌ Error al verificar/crear colección: {e}", exc_info=True)
+            raise
+    
+    def load_data(self, parquet_path: str) -> bool:
+        """
+        Carga los embeddings desde un parquet a Qdrant
+        
+        Args:
+            parquet_path: Ruta al archivo parquet con embeddings
+            
+        Returns:
+            True si la carga fue exitosa, False si falló
+        """
+        try:
+            # Leer parquet
+            df = pd.read_parquet(parquet_path)
+            
+            if "embedding" not in df.columns:
+                raise ValueError(f"El parquet {parquet_path} no contiene columna 'embedding'")
+            
+            # Preparar puntos para Qdrant
+            points: List[PointStruct] = []
+            
+            for idx, row in df.iterrows():
+                # Generar ID único
+                point_id = str(uuid.uuid4())
+                
+                # Preparar payload (todos los campos excepto embedding)
+                payload = {col: row[col] for col in df.columns if col != "embedding"}
+                
+                # Crear punto
+                point = PointStruct(
+                    id=point_id,
+                    vector=row["embedding"],
+                    payload=payload
+                )
+                points.append(point)
+            
+            # Insertar en Qdrant (en batch)
+            self.client.upsert(
+                collection_name=self.collection_name,
+                points=points
+            )
+            
+            logger.info(f"✅ {len(points)} embeddings cargados a Qdrant desde {parquet_path}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Error al cargar embeddings a Qdrant desde {parquet_path}: {e}", exc_info=True)
+            return False

transform/transform.py

@@ -0,0 +1,42 @@
+from contracts.contracts import transform_data
+from pathlib import Path
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+import logging
+from typing import Optional
+
+# Configurar logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+class transform_embedding(transform_data):
+    def transform_data(self, parquet_path: str) -> Optional[str]:
+        try:
+            # Cargar datos del parquet
+            df = pd.read_parquet(parquet_path)
+            if "text" not in df.columns:
+                raise ValueError("El parquet no contiene una columna 'text' para generar embeddings.")
+
+            # Inicializar modelo con embeddings de 768 dimensiones
+            model = SentenceTransformer("all-mpnet-base-v2")
+
+            # Generar embeddings
+            embeddings = model.encode(df["text"].tolist(), show_progress_bar=True)
+            df["embedding"] = embeddings.tolist()  # guardar como lista
+
+            # Guardar parquet procesado
+            output_dir = Path("datos_embeddings")
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+            output_file = output_dir / Path(parquet_path).name
+            df.to_parquet(output_file, index=False)
+
+            logger.info(f"✅ Embeddings (768 dim) generados y guardados en: {output_file}")
+            return str(output_file)
+
+        except Exception as e:
+            logger.error(f"❌ Error al transformar en embeddings {parquet_path}: {e}", exc_info=True)
+            return None