langchain/libs/experimental/langchain_experimental/open_clip/open_clip.py
Eugene Yurtsev 86355640c3
experimental[patch]: Use get_fields adapter (#25193)
Change all usages of __fields__ with get_fields adapter merged into
langchain_core.

Code mod generated using the following grit pattern:

```
engine marzano(0.1)
language python


`$X.__fields__` => `get_fields($X)` where {
    add_import(source="langchain_core.utils.pydantic", name="get_fields")
}
```
2024-08-08 15:10:11 -04:00

93 lines
3.4 KiB
Python

from typing import Any, Dict, List
from langchain.pydantic_v1 import BaseModel, root_validator
from langchain_core.embeddings import Embeddings
from langchain_core.utils.pydantic import get_fields
class OpenCLIPEmbeddings(BaseModel, Embeddings):
"""OpenCLIP Embeddings model."""
model: Any
preprocess: Any
tokenizer: Any
# Select model: https://github.com/mlfoundations/open_clip
model_name: str = "ViT-H-14"
checkpoint: str = "laion2b_s32b_b79k"
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
"""Validate that open_clip and torch libraries are installed."""
try:
import open_clip
# Fall back to class defaults if not provided
model_name = values.get("model_name", get_fields(cls)["model_name"].default)
checkpoint = values.get("checkpoint", get_fields(cls)["checkpoint"].default)
# Load model
model, _, preprocess = open_clip.create_model_and_transforms(
model_name=model_name, pretrained=checkpoint
)
tokenizer = open_clip.get_tokenizer(model_name)
values["model"] = model
values["preprocess"] = preprocess
values["tokenizer"] = tokenizer
except ImportError:
raise ImportError(
"Please ensure both open_clip and torch libraries are installed. "
"pip install open_clip_torch torch"
)
return values
def embed_documents(self, texts: List[str]) -> List[List[float]]:
text_features = []
for text in texts:
# Tokenize the text
tokenized_text = self.tokenizer(text)
# Encode the text to get the embeddings
embeddings_tensor = self.model.encode_text(tokenized_text)
# Normalize the embeddings
norm = embeddings_tensor.norm(p=2, dim=1, keepdim=True)
normalized_embeddings_tensor = embeddings_tensor.div(norm)
# Convert normalized tensor to list and add to the text_features list
embeddings_list = normalized_embeddings_tensor.squeeze(0).tolist()
text_features.append(embeddings_list)
return text_features
def embed_query(self, text: str) -> List[float]:
return self.embed_documents([text])[0]
def embed_image(self, uris: List[str]) -> List[List[float]]:
try:
from PIL import Image as _PILImage
except ImportError:
raise ImportError("Please install the PIL library: pip install pillow")
# Open images directly as PIL images
pil_images = [_PILImage.open(uri) for uri in uris]
image_features = []
for pil_image in pil_images:
# Preprocess the image for the model
preprocessed_image = self.preprocess(pil_image).unsqueeze(0)
# Encode the image to get the embeddings
embeddings_tensor = self.model.encode_image(preprocessed_image)
# Normalize the embeddings tensor
norm = embeddings_tensor.norm(p=2, dim=1, keepdim=True)
normalized_embeddings_tensor = embeddings_tensor.div(norm)
# Convert tensor to list and add to the image_features list
embeddings_list = normalized_embeddings_tensor.squeeze(0).tolist()
image_features.append(embeddings_list)
return image_features