How to Use LLM2Vec for Text Encoding

Apr 12, 2024 | Educational

In the era of natural language processing, Large Language Models (LLMs) have proven to be powerful tools for a variety of text tasks. One such tool is LLM2Vec, which transforms decoder-only LLMs into effective text encoders. This article will guide you through the steps to effectively employ LLM2Vec, complete with troubleshooting tips!

Installation

Getting started with LLM2Vec is straightforward. You need to install it via pip. Open your terminal and run:

bash
pip install llm2vec

Usage

Let’s dive into the steps to utilize LLM2Vec in your Python projects.

Step 1: Import Necessary Libraries

You’ll start by importing the essential libraries and modules.

python
from llm2vec import LLM2Vec
import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig
from peft import PeftModel

Step 2: Load the Base Model

Think of this step as setting up your kitchen before cooking. You need to gather all of your ingredients (model components) to create a delicious dish (effective text encoding). Here’s how you do it:

python
tokenizer = AutoTokenizer.from_pretrained("McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp")
config = AutoConfig.from_pretrained("McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp",
    trust_remote_code=True,
    config=config,
    torch_dtype=torch.bfloat16,
    device_map="cuda" if torch.cuda.is_available() else "cpu"
)
model = PeftModel.from_pretrained(model, "McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp")
model = model.merge_and_unload()  # This can take several minutes on CPU

Step 3: Load the Unsupervised SimCSE Model

Now, let’s load an unsupervised SimCSE model which serves as a spice to enhance the flavor of your dish. This integrates additional trained weights to the base model.

python
model = PeftModel.from_pretrained(model, "McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp-unsup-simcse")
l2v = LLM2Vec(model, tokenizer, pooling_mode="mean", max_length=512)

Step 4: Encoding Queries and Documents

With all the ingredients in place, you can proceed to cook. Here’s how to encode queries and documents:

python
instruction = "Given a web search query, retrieve relevant passages that answer the query:"
queries = [
    [instruction, "how much protein should a female eat"],
    [instruction, "summit define"],
]
q_reps = l2v.encode(queries)

documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day.",
    "Definition of summit for English Language Learners: 1  the highest point of a mountain.",
]
d_reps = l2v.encode(documents)

Step 5: Compute Cosine Similarity

Now it’s time to see how your dish tastes! You will compute the cosine similarity between the encoded queries and documents:

python
q_reps_norm = torch.nn.functional.normalize(q_reps, p=2, dim=1)
d_reps_norm = torch.nn.functional.normalize(d_reps, p=2, dim=1)
cos_sim = torch.mm(q_reps_norm, d_reps_norm.transpose(0, 1))
print(cos_sim)

Troubleshooting

If you encounter any issues during the installation or while running the code, here are some common troubleshooting tips:

  • Error: “ModuleNotFoundError” – This usually means that a library is not installed. Double-check your pip installations.
  • Performance Issues: If the encoding process is slow, ensure that your GPU is being utilized by checking your CUDA settings.
  • Cosine Similarity Matrix Is Empty: This could be due to improperly formatted queries or documents. Ensure that all inputs are in the expected format.

For more insights, updates, or to collaborate on AI development projects, stay connected with fxis.ai.

Conclusion

LLM2Vec opens doors to new dimensions in text encoding using powerful existing models. With just a few steps, you can leverage this technique for various natural language processing tasks.

At fxis.ai, we believe that such advancements are crucial for the future of AI, as they enable more comprehensive and effective solutions. Our team is continually exploring new methodologies to push the envelope in artificial intelligence, ensuring that our clients benefit from the latest technological innovations.

Stay Informed with the Newest F(x) Insights and Blogs

Tech News and Blog Highlights, Straight to Your Inbox