How to Use the GTE-Small Model for Text Embedding

Mar 19, 2024 | Educational

The GTE-Small model developed by Alibaba DAMO Academy is a powerful tool for handling various downstream tasks in natural language processing, such as information retrieval and semantic textual similarity. In this guide, we will walk you through how to utilize this model effectively using both Python and JavaScript.

Understanding the GTE-Small Model

The GTE models are akin to seasoned chefs crafting gourmet dishes. Just as chefs use different ingredients to create a diverse menu, GTE models utilize a large-scale corpus of text pairs. Each model size—from GTE-Large to GTE-Base—is tailored for specific tasks and resource capabilities. The GTE-Small offers a compact solution without compromising the performance needed for core tasks in text embedding. Essentially, it’s your go-to choice for a quick and efficient meal in the world of embeddings!

Using GTE-Small Model in Python

To incorporate GTE-Small using Python, follow these simple steps:

  1. Make sure you have the required libraries installed: Transformers and PyTorch.
  2. Use the following code snippet to generate embeddings:
  3. import torch.nn.functional as F
    from torch import Tensor
    from transformers import AutoTokenizer, AutoModel
    
    def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
    
    input_texts = [
        "what is the capital of China?",
        "how to implement quick sort in python?",
        "Beijing",
        "sorting algorithms"
    ]
    
    tokenizer = AutoTokenizer.from_pretrained("Supabase/gte-small")
    model = AutoModel.from_pretrained("Supabase/gte-small")
    
    batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors="pt")
    outputs = model(**batch_dict)
    embeddings = average_pool(outputs.last_hidden_state, batch_dict["attention_mask"])
    embeddings = F.normalize(embeddings, p=2, dim=1)
    
    scores = (embeddings[:1] @ embeddings[1:].T) * 100
    print(scores.tolist())
  4. This code will generate clustered embeddings that you can use for various applications!

Using GTE-Small Model in JavaScript

For JavaScript enthusiasts, here’s how to use the GTE-Small model with Transformers.js:

  1. Set up your environment with Deno or Node.js.
  2. Use the snippet below for extracting embeddings:
  3. import { serve } from "https://deno.land/std@0.168.0/http/server.ts";
    import { env, pipeline } from "https://cdn.jsdelivr.net/npm/@xenova/transformers@2.5.0";
    
    const pipe = await pipeline("feature-extraction", "Supabase/gte-small");
    
    serve(async (req) => {
        const input = await req.json();
        const output = await pipe(input, { pooling: 'mean', normalize: true });
        const embedding = Array.from(output.data);
        return new Response(JSON.stringify(embedding), { headers: { "Content-Type": "application/json" } });
    });
  4. Now, you can handle input strings through your server and generate embeddings seamlessly!

Troubleshooting Tips

While using the GTE-Small model, you may encounter a few challenges. Here are some troubleshooting tips:

  • If your input text is truncated, remember that the model only handles a maximum of 512 tokens. This is like feeding a little bird, where you shouldn’t overstuff it!
  • Make sure your environment is set up correctly with all necessary libraries and dependencies.
  • Feel free to refer to the MTEB leaderboard for performance metrics to choose the right model for your needs.
  • If issues persist, document any error messages and check community forums or resources online.

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

Conclusion

Using the GTE-Small model opens up a world of possibilities for natural language processing tasks. With Python or JavaScript, you can tap into the power of these embeddings effectively. 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