Spaces:
Sleeping
Sleeping
File size: 1,553 Bytes
b5739f3 |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import os
from transformers import AutoModel, AutoTokenizer
import torch
from dotenv import load_dotenv
from langchain.schema.embeddings import Embeddings
load_dotenv() # ✅ make sure .env is read
class GemmaEmbeddings:
def __init__(self, model_name="google/embeddinggemma-300m", device=None):
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
hf_token = os.environ.get("HUGGINGFACETOEN")
if not hf_token:
raise ValueError("❌ Hugging Face token not found. Please set HF_TOKEN in .env")
# ✅ Pass token when loading
self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
self.model = AutoModel.from_pretrained(model_name, use_auth_token=hf_token).to(self.device)
def embed(self, texts):
if isinstance(texts, str):
texts = [texts]
encodings = self.tokenizer(
texts, padding=True, truncation=True, return_tensors="pt"
).to(self.device)
with torch.no_grad():
model_output = self.model(**encodings)
embeddings = model_output.last_hidden_state.mean(dim=1).cpu().numpy()
return embeddings.tolist()
class GemmaLangChainEmbeddings(Embeddings):
def __init__(self, model_name="google/embeddinggemma-300m"):
self.gemma = GemmaEmbeddings(model_name=model_name)
def embed_query(self, text: str):
return self.gemma.embed(text)[0]
def embed_documents(self, texts: list[str]):
return self.gemma.embed(texts)
|