Is it possible to use llms other than openai?

1

I'd like to know if I can use gemini for llm annotation.

2

Welcome to the forum @hazelkang :waving_hand:

Yes, absolutely! You can integrate any LLM, including Google's Gemini models, into your Prodigy annotation workflow. The key is to use the spacy-llm library to create a custom backend that communicates with the Gemini API.

Here’s a how to set it up for a recipe like ner.llm.correct. There are 3 main tasks:

  1. Create a custom spacy-llm model class: Write a Python script that defines how to call the Gemini API and registers it with spacy-llm
  2. Create a spacy-llm config file: This will configure the spaCy pipeline used in Prodigy to use your Gemini backend.
  3. Run the Prodigy Recipe: Launch Prodigy with the correct command-line arguments to load your custom code and configuration.

Re 1. Gemini spacy-llm backend

First, you need a Python script that contains the logic for interacting with the Gemini API. This script defines a custom class and registers it with spacy-llm so Prodigy can find it. You can check how this is done for the built-in spacy-llm models here. We'll basically follow the same pattern (only, instead of low lever requests you could use generativeai helper method) :

import os
import warnings
from typing import Any, Callable, Dict, Iterable, List, Optional

import google.generativeai as genai
from confection import SimpleFrozenDict
from spacy_llm.registry import registry
from spacy_llm.models.rest.base import REST

# --- Part 1: The Backend Class ---


class Gemini(REST):
    """Backend for the Gemini API."""

    def __init__(
        self,
        name: str,
        config: Dict[str, Any],
        strict: bool,
        max_tries: int,
        interval: float,
        max_request_time: float,
        context_length: Optional[int] = None,
    ):
        super().__init__(
            name=name,
            endpoint="",
            config=config,
            strict=strict,
            max_tries=max_tries,
            interval=interval,
            max_request_time=max_request_time,
            context_length=context_length,
        )
        genai.configure(api_key=self.credentials["api_key"])
        self._model = genai.GenerativeModel(self._name)

    @property
    def credentials(self) -> Dict[str, str]:
        api_key = os.getenv("GEMINI_API_KEY")
        if api_key is None:
            warnings.warn("Could not find GEMINI_API_KEY env var")
        assert api_key is not None
        return {"api_key": api_key}

    def _verify_auth(self) -> None:
        try:
            genai.list_models()
        except Exception as e:
            raise ValueError(f"Could not authenticate with GEMINI_API_KEY: {e}") from e

    @staticmethod
    def _get_context_lengths() -> Dict[str, int]:
        return {"gemini-1.5-flash": 1048576}

    def __call__(self, prompts: Iterable[Iterable[str]]) -> Iterable[Iterable[str]]:
        all_api_responses: List[List[str]] = []
        for prompts_for_doc in prompts:
            api_responses_for_doc: List[str] = []
            for prompt in prompts_for_doc:
                try:
                    api_response = self._model.generate_content(
                        prompt,
                        generation_config=self._config,
                    )
                    full_response_text = (
                        "".join(part.text for part in api_response.parts)
                        if api_response.parts
                        else ""
                    )
                    api_responses_for_doc.append(full_response_text.strip())
                except Exception as e:
                    if self._strict:
                        raise e
                    warnings.warn(f"API call failed: {e}")
                    api_responses_for_doc.append("")
            all_api_responses.append(api_responses_for_doc)
        return all_api_responses


# --- Part 2: The Registry Function ---


@registry.llm_models("spacy.Gemini.v1")
def gemini_v1(
    config: Dict[str, Any] = SimpleFrozenDict(),
    name: str = "gemini-1.5-flash",
    strict: bool = False,
    max_tries: int = 5,
    interval: float = 1.0,
    max_request_time: float = 60,
    context_length: Optional[int] = None,
) -> Callable[[Iterable[Iterable[str]]], Iterable[Iterable[str]]]:
    """Registers the Gemini model integration."""
    return Gemini(
        name=name,
        config=config,
        strict=strict,
        max_tries=max_tries,
        interval=interval,
        max_request_time=max_request_time,
        context_length=context_length,
    )

RE 2 spacy-llm config

Next, create a gemini.cfg file. This tells spaCy how to configure the spacy-llm pipeline that will generate the initial NER suggestions using your new Gemini backend. For example:

[nlp]
lang = "en"
pipeline = ["llm"]

[components]
[components.llm]
factory = "llm"
save_io = true

[components.llm.task]
@llm_tasks = "spacy.NER.v3"
labels = ["PERSON", "ORGANIZATION", "PRODUCT"]

[components.llm.model]
@llm_models = "spacy.Gemini.v1"
name = "gemini-1.5-flash"
config = {"temperature": 0.0}

RE 3 Run Prodigy recipe importing the custom spacy-llm code and the config
Assuming you have GEMINI_API_KEY in your .envfile you could run:

dotenv run -- python -m prodigy ner.llm.correct test gemini.cfg input.jsonl -F custom_gemini.py

Hope this NER example is helpful to get you started on your task!

1 Like
3

Thanks! Also wondering if I buy the one private plan, can I work with other RAs on one project? or do I have to buy multiple accounts to work with one? Since for my annotation task it is important to cross-verify with other collaboraters.

4

Hi @hazelkang,

If the other RAs will only access the hosted web app (i.e. they would be acting as annotators) then that's totally fine under the Personal license. There's no limit on the number of annotators.
A license is required for everyone who has access to the Prodigy Python library and back-end.

1 Like