Spaces:
Sleeping
Sleeping
File size: 1,069 Bytes
6c975bc 4c932d5 995f7ff 16b4fa5 995f7ff 4c932d5 6c975bc 995f7ff 6c975bc afd159e 7a57ef8 0296f81 995f7ff 6c975bc e3e85bf afd159e 7a57ef8 0296f81 995f7ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
from transformers import AutoModel
from PIL import Image
import requests
from io import BytesIO
import os
import torch
import numpy as np
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
os.environ['HF_HOME'] = '/app/hf_cache'
# Load model
model = AutoModel.from_pretrained('jinaai/jina-clip-v2', trust_remote_code=True).to(device)
def get_text_embedding(texts, truncate_dim=512):
embeddings = model.encode_text(texts, truncate_dim=truncate_dim)
# if isinstance(embeddings, np.ndarray):
embeddings = torch.from_numpy(embeddings)
print(embeddings)
return embeddings
def get_image_embedding(image_urls, truncate_dim=512):
"""
Takes a list of image URLs and returns embeddings using model.encode_image.
Assumes model.encode_image supports URL input directly.
"""
embeddings = model.encode_image(image_urls, truncate_dim=truncate_dim)
# if not isinstance(embeddings, torch.Tensor):
embeddings = torch.tensor(embeddings)
print(embeddings)
return embeddings
|