Granite Embedding

pip install flash_attn

pip install sentence_transformers

from sentence_transformers import SentenceTransformer, util

model_path = "ibm-granite/granite-embedding-311m-multilingual-r2"
# Load the Sentence Transformer model
model = SentenceTransformer(model_path)

input_queries = [
    'What is the tallest mountain in Japan?',          # English query
    'Wer hat das Lied Achy Breaky Heart geschrieben?', # German query
    'ドイツの首都はどこですか？',                            # Japanese query
    ]

input_passages = [
    "富士山は、静岡県と山梨県にまたがる活火山で、標高3776.12 mで日本最高峰の独立峰である。",  # Japanese passage
    "Achy Breaky Heart is a country song written by Don Von Tress. Originally titled Don't Tell My Heart and performed by The Marcy Brothers in 1991.",  # English passage
    "Berlin ist die Hauptstadt und ein Land der Bundesrepublik Deutschland. Die Stadt ist mit rund 3,7 Millionen Einwohnern die bevölkerungsreichste Kommune Deutschlands.",  # German passage
    ]

# Cross-lingual retrieval: each query should score highest with its matching passage in a different language
query_embeddings = model.encode(input_queries)
passage_embeddings = model.encode(input_passages)

# calculate cosine similarity — expect high scores on the diagonal (EN→JA, DE→EN, JA→DE)
print(util.cos_sim(query_embeddings, passage_embeddings))
# output: tensor([[0.9393, 0.6899, 0.7627],
#                 [0.6780, 0.9598, 0.7062],
#                 [0.7818, 0.7342, 0.9172]])

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("ibm-granite/granite-embedding-311m-multilingual-r2")

# Full 768-dimensional embeddings
full_embeddings = model.encode(["example text"])
print(full_embeddings.shape)  # (1, 768)

# Truncated to 256 dimensions — still effective for many retrieval tasks
truncated_embeddings = model.encode(["example text"], truncate_dim=256)
print(truncated_embeddings.shape)  # (1, 256)

pip install transformers torch

import torch
from transformers import AutoModel, AutoTokenizer

model_path = "ibm-granite/granite-embedding-311m-multilingual-r2"

# Load the model and tokenizer
model = AutoModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.eval()

input_queries = [
    'What is the tallest mountain in Japan?',          # English query
    'Wer hat das Lied Achy Breaky Heart geschrieben?', # German query
    'ドイツの首都はどこですか？',                            # Japanese query
    ]

# tokenize inputs
tokenized_queries = tokenizer(input_queries, padding=True, truncation=True, return_tensors='pt')

# encode queries
with torch.no_grad():
    model_output = model(**tokenized_queries)
    # Perform pooling. granite-embedding-311m-multilingual-r2 uses CLS Pooling
    query_embeddings = model_output[0][:, 0]

# normalize the embeddings
query_embeddings = torch.nn.functional.normalize(query_embeddings, dim=1)

pip install langchain_huggingface sentence_transformers \
    langchain_milvus 'pymilvus[milvus_lite]' \
    langchain_community \
    langchain_text_splitters \
    wget

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_milvus import Milvus
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import CharacterTextSplitter
import os, tempfile, wget

# load the embedding model
embeddings_model = HuggingFaceEmbeddings(model_name="ibm-granite/granite-embedding-30m-english")

# setup the vectordb
db_file = tempfile.NamedTemporaryFile(prefix="milvus_", suffix=".db", delete=False).name
print(f"The vector database will be saved to {db_file}")
vector_db = Milvus(
  embedding_function=embeddings_model,
  connection_args={"uri": db_file},
  auto_id=True,
  index_params={"index_type": "AUTOINDEX"},
  )

# load example corpus file
filename = 'state_of_the_union.txt'
url = 'https://raw.github.com/IBM/watson-machine-learning-samples/master/cloud/data/foundation_models/state_of_the_union.txt'

if not os.path.isfile(filename):
  wget.download(url, out=filename)

loader = TextLoader(filename)
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=0)
texts = text_splitter.split_documents(documents)

# add processed documents to the vectordb
vector_db.add_documents(texts)

# search the vectordb with the query
query = "What did the president say about Ketanji Brown Jackson"
docs = vector_db.similarity_search(query)
print(docs[0].page_content)