Add embedding service and preprocessor; integrate with Gradio interface

- Introduced `EmbeddingService` for generating text embeddings using sentence-transformers.
- Added `TextPreprocessor` for cleaning and tokenizing input text.
- Created a new endpoint for obtaining embeddings and integrated it into the Gradio interface.
- Updated `requirements.txt` to include necessary libraries.
- Added configuration settings in `config.py` for model and server parameters.

app.py

@@ -1,11 +1,17 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+from app.services.embedding_service import EmbeddingService
+from app.config import EMBEDDING_MODEL  # Import from config
+from app.services.preprocessor import TextPreprocessor
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+# Initialize EmbeddingService
+embedding_service = EmbeddingService(model_name=EMBEDDING_MODEL, preprocessor=TextPreprocessor())
+
 
 def respond(
     message,
@@ -39,6 +45,26 @@ def respond(
         response += token
         yield response
 
+def get_embedding(text: str) -> list[float]:
+    """
+    Endpoint to get the embedding of a text.
+    """
+    try:
+        return embedding_service.get_embedding(text)
+    except ValueError as e:
+        # Handle the case where the input text is too long
+        return f"Error: {str(e)}"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create a separate Gradio interface for the embedding endpoint
+embedding_iface = gr.Interface(
+    fn=get_embedding,
+    inputs=gr.Textbox(placeholder="Enter text here...", label="Input Text"),
+    outputs=gr.JSON(label="Embedding"),  # Use JSON output for the embedding vector
+    title="Embedding Service",
+    description="Get the embedding of a text using the Vietnamese Bi-Encoder.",
+)
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -59,6 +85,9 @@ demo = gr.ChatInterface(
     ],
 )
 
+# Combine the interfaces
+demo = gr.TabbedInterface([demo, embedding_iface], ["Chatbot", "Embedding"])
+
 
 if __name__ == "__main__":
     demo.launch()

app/config.py

@@ -0,0 +1,25 @@
+import os
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv(override=True)
+
+# Application settings
+APP_NAME = "Vietnamese RAG"
+DEBUG = os.getenv("DEBUG", "False").lower() in ("true", "1", "t")
+API_PREFIX = "/api"
+
+# Model settings
+EMBEDDING_MODEL = "bkai-foundation-models/vietnamese-bi-encoder"
+MAX_TOKEN_LIMIT = 128
+DEFAULT_CHUNK_SIZE = 110  # Safe margin below MAX_TOKEN_LIMIT
+DEFAULT_CHUNK_OVERLAP = 20
+DEFAULT_TOP_K = 5
+
+# Server settings
+HOST = os.getenv("HOST", "0.0.0.0")
+PORT = int(os.getenv("PORT", "8000"))
+
+# Cache settings
+ENABLE_CACHE = os.getenv("ENABLE_CACHE", "True").lower() in ("true", "1", "t")
+CACHE_SIZE = int(os.getenv("CACHE_SIZE", "1000")) 

app/services/embedding_service.py

@@ -0,0 +1,194 @@
+import logging
+from typing import List, Dict, Any, Optional, Union
+import numpy as np
+from functools import lru_cache
+from sentence_transformers import SentenceTransformer
+
+from app.config import (
+    EMBEDDING_MODEL, 
+    MAX_TOKEN_LIMIT, 
+    ENABLE_CACHE, 
+    CACHE_SIZE
+)
+from app.services.preprocessor import TextPreprocessor
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingService:
+    """Service for generating embeddings for text using sentence-transformers."""
+    
+    def __init__(self, model_name: str = EMBEDDING_MODEL, preprocessor=None):
+        """
+        Initialize the embedding service.
+        
+        Args:
+            model_name: Name of the sentence-transformers model to use
+            preprocessor: Optional TextPreprocessor instance
+        """
+        logger.info(f"Loading embedding model: {model_name}")
+        self.model = SentenceTransformer(model_name)
+        self.model_dim = self.model.get_sentence_embedding_dimension()
+        logger.info(f"Model loaded. Embedding dimension: {self.model_dim}")
+        
+        # Use provided preprocessor or create one
+        self.preprocessor = preprocessor or TextPreprocessor()
+        
+        # Set up caching if enabled
+        if ENABLE_CACHE:
+            self.get_embedding = lru_cache(maxsize=CACHE_SIZE)(self._get_embedding)
+        else:
+            self.get_embedding = self._get_embedding
+    
+    def _get_embedding(self, text: str) -> List[float]:
+        """
+        Generate embedding for a text string.
+        
+        Args:
+            text: Text to generate embedding for
+            
+        Returns:
+            List of floats representing the embedding vector
+        """
+        if not text or not isinstance(text, str):
+            logger.warning("Empty or invalid text provided for embedding generation")
+            return [0.0] * self.model_dim
+        
+        # Use preprocessor for token counting only
+        token_count = self.preprocessor.count_tokens(text)
+        
+        # Check against token limit
+        if token_count > MAX_TOKEN_LIMIT:
+            logger.error(
+                f"Text exceeds max token limit ({token_count} > {MAX_TOKEN_LIMIT}). "
+                f"Please chunk your text before encoding."
+            )
+            raise ValueError(f"Text exceeds max token limit ({token_count} > {MAX_TOKEN_LIMIT})")
+            
+        try:
+            # Directly encode the text string
+            embedding = self.model.encode(text).tolist()
+            return embedding
+        except Exception as e:
+            logger.error(f"Error generating embedding: {str(e)}")
+            return [0.0] * self.model_dim
+    
+    def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate embeddings for a batch of texts.
+        
+        Args:
+            texts: List of texts to generate embeddings for
+            
+        Returns:
+            List of embedding vectors
+        """
+        if not texts:
+            return []
+        
+        # Validate texts are within token limit
+        for i, text in enumerate(texts):
+            if not text or not isinstance(text, str):
+                logger.warning(f"Empty or invalid text at index {i}")
+                continue
+                
+            # Check token count
+            token_count = self.preprocessor.count_tokens(text)
+            if token_count > MAX_TOKEN_LIMIT:
+                logger.error(
+                    f"Text at index {i} exceeds max token limit ({token_count} > {MAX_TOKEN_LIMIT}). "
+                    f"Please chunk your text before encoding."
+                )
+                raise ValueError(f"Text at index {i} exceeds max token limit ({token_count} > {MAX_TOKEN_LIMIT})")
+        
+        try:
+            # Let the model handle the batch encoding directly
+            embeddings = self.model.encode(texts).tolist()
+            return embeddings
+        except Exception as e:
+            logger.error(f"Error generating batch embeddings: {str(e)}")
+            return [[0.0] * self.model_dim] * len(texts)
+    
+    def embed_chunks(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Generate embeddings for a list of text chunks.
+        
+        Args:
+            chunks: List of chunk dictionaries with text and metadata
+            
+        Returns:
+            List of chunk dictionaries with added embeddings
+        """
+        if not chunks:
+            return []
+        
+        # Extract texts from chunks
+        texts = [chunk["text"] for chunk in chunks]
+        
+        # Generate embeddings
+        embeddings = self.get_embeddings_batch(texts)
+        
+        # Add embeddings to chunks
+        result_chunks = []
+        for chunk, embedding in zip(chunks, embeddings):
+            chunk_with_embedding = chunk.copy()
+            chunk_with_embedding["embedding"] = embedding
+            result_chunks.append(chunk_with_embedding)
+        
+        return result_chunks
+    
+    def similarity_search(
+        self, 
+        query: str, 
+        embeddings: List[List[float]], 
+        texts: List[str],
+        metadata: Optional[List[Dict[str, Any]]] = None,
+        top_k: int = 5
+    ) -> List[Dict[str, Any]]:
+        """
+        Find the most similar texts to a query.
+        
+        Args:
+            query: Query text
+            embeddings: List of embedding vectors to search
+            texts: List of texts corresponding to the embeddings
+            metadata: Optional list of metadata for each text
+            top_k: Number of top matches to return
+            
+        Returns:
+            List of matches with text, score, and metadata
+        """
+        if not query or not embeddings or not texts:
+            return []
+        
+        if metadata is None:
+            metadata = [{} for _ in range(len(texts))]
+        
+        # Generate query embedding
+        query_embedding = self.get_embedding(query)
+        
+        # Convert to numpy arrays for efficient computation
+        query_embedding_np = np.array(query_embedding)
+        embeddings_np = np.array(embeddings)
+        
+        # Compute cosine similarity
+        similarity_scores = np.dot(embeddings_np, query_embedding_np) / (
+            np.linalg.norm(embeddings_np, axis=1) * np.linalg.norm(query_embedding_np)
+        )
+        
+        # Get top-k indices
+        if top_k > len(texts):
+            top_k = len(texts)
+            
+        top_indices = np.argsort(similarity_scores)[-top_k:][::-1]
+        
+        # Prepare results
+        results = []
+        for idx in top_indices:
+            results.append({
+                "text": texts[idx],
+                "score": float(similarity_scores[idx]),
+                "metadata": metadata[idx]
+            })
+        
+        return results 

app/services/preprocessor.py

@@ -0,0 +1,38 @@
+import re
+from transformers import AutoTokenizer
+from app.config import EMBEDDING_MODEL
+
+class TextPreprocessor:
+    """
+    A simple text preprocessor for cleaning and tokenizing text.
+    """
+
+    def __init__(self, model_name: str = EMBEDDING_MODEL):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def clean_text(self, text: str) -> str:
+        """
+        Remove extra whitespace and control characters from text.
+
+        Args:
+            text: The text to clean.
+
+        Returns:
+            The cleaned text.
+        """
+        text = re.sub(r"[\s\t\n]+", " ", text)  # Normalize whitespace
+        text = re.sub(r"[\x00-\x1F\x7F]", "", text)  # Remove control characters
+        return text.strip()
+
+    def count_tokens(self, text: str) -> int:
+        """
+        Count the number of tokens in the text using a tokenizer.
+
+        Args:
+            text: The text to tokenize.
+
+        Returns:
+            The number of tokens.
+        """
+        # Tokenize the text and return the length of the input IDs
+        return len(self.tokenizer(text).input_ids) 

requirements.txt

@@ -1 +1,5 @@
-huggingface_hub==0.25.2
+huggingface_hub==0.25.2
+gradio 
+sentence-transformers 
+python-dotenv 
+transformers