111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
"""
|
|
Embeddings utilities using Ollama and Nomic.
|
|
"""
|
|
import os
|
|
import requests
|
|
import numpy as np
|
|
from typing import List, Optional, Union
|
|
import ollama
|
|
from app.config import Config
|
|
|
|
class EmbeddingService:
|
|
"""Service for generating embeddings using Ollama and Nomic."""
|
|
|
|
@staticmethod
|
|
def get_ollama_embeddings(texts: List[str], model: Optional[str] = None) -> List[List[float]]:
|
|
"""
|
|
Generate embeddings using Ollama.
|
|
|
|
Args:
|
|
texts: List of texts to generate embeddings for
|
|
model: Ollama model to use for embeddings (default from config)
|
|
|
|
Returns:
|
|
List of embeddings as float arrays
|
|
"""
|
|
if model is None:
|
|
# Use model from config
|
|
model = Config.OLLAMA_MODEL
|
|
|
|
# Set Ollama host from config
|
|
ollama.host = Config.OLLAMA_HOST
|
|
|
|
embeddings = []
|
|
|
|
for text in texts:
|
|
try:
|
|
# Call Ollama API for embeddings
|
|
response = ollama.embeddings(model=model, prompt=text)
|
|
embedding = response.get("embedding", [])
|
|
embeddings.append(embedding)
|
|
except Exception as e:
|
|
print(f"Error generating Ollama embedding: {e}")
|
|
# Return a zero embedding as fallback
|
|
embeddings.append([0.0] * 768) # typical dimension for text embeddings
|
|
|
|
return embeddings
|
|
|
|
@staticmethod
|
|
def get_nomic_embeddings(texts: List[str]) -> List[List[float]]:
|
|
"""
|
|
Generate embeddings using Nomic.
|
|
|
|
Args:
|
|
texts: List of texts to generate embeddings for
|
|
|
|
Returns:
|
|
List of embeddings as float arrays
|
|
"""
|
|
try:
|
|
# The new version of Nomic requires a Cohere API key, so we'll fall back to Ollama
|
|
# if we don't have one configured
|
|
cohere_api_key = Config.COHERE_API_KEY
|
|
|
|
if not cohere_api_key:
|
|
print("No Cohere API key found for Nomic embeddings, falling back to Ollama")
|
|
return EmbeddingService.get_ollama_embeddings(texts)
|
|
|
|
# Dynamically import nomic embedders to avoid startup errors if not available
|
|
from nomic.embedders import CohereEmbedder
|
|
|
|
# Create a Nomic embedding model using CohereEmbedder with API key
|
|
embedding_model = CohereEmbedder(cohere_api_key=cohere_api_key)
|
|
|
|
# Generate embeddings for the texts
|
|
embeddings = []
|
|
for text in texts:
|
|
embedding = embedding_model.embed(text)
|
|
embeddings.append(embedding)
|
|
|
|
return embeddings
|
|
except Exception as e:
|
|
print(f"Error generating Nomic embeddings: {e}")
|
|
# Fall back to Ollama embeddings
|
|
print("Falling back to Ollama embeddings")
|
|
return EmbeddingService.get_ollama_embeddings(texts)
|
|
|
|
@staticmethod
|
|
def get_embeddings(texts: Union[str, List[str]], use_nomic: Optional[bool] = None) -> List[List[float]]:
|
|
"""
|
|
Generate embeddings using either Nomic or Ollama.
|
|
|
|
Args:
|
|
texts: Text or list of texts to generate embeddings for
|
|
use_nomic: Whether to use Nomic (True) or Ollama (False), defaults to config setting
|
|
|
|
Returns:
|
|
List of embeddings as float arrays
|
|
"""
|
|
# Convert single text to list
|
|
if isinstance(texts, str):
|
|
texts = [texts]
|
|
|
|
# If use_nomic is not specified, use the config setting
|
|
if use_nomic is None:
|
|
use_nomic = Config.USE_NOMIC_EMBEDDINGS
|
|
|
|
# Generate embeddings using chosen method
|
|
if use_nomic:
|
|
return EmbeddingService.get_nomic_embeddings(texts)
|
|
else:
|
|
return EmbeddingService.get_ollama_embeddings(texts) |