Extract multiple entities

[kormilitzin]

Hi all, I am using local LLMs to extract certain information from text, however the provided examples can only pull our only the first mention of the required information.

For example, you have a text: "I like travelling and visited Berlin, Paris and London." If I run the model, I get only "Berlin" identified and all other cities are ignored.

Below is my code snippet to reproduce the error. I am using a local instance of LLMs via Ollama.

from openai import OpenAI
from pydantic import BaseModel, Field
import instructor

class Cities(BaseModel):
    city: str = Field(
        description="A city that the user visited."
    )
    class Config:
        arbitrary_types_allowed = True


text = "I like travelling and visited Berlin, Paris and London."

query = f"""
You are an information extraction tool. The following is a text:
{text}

Please extract all mentions of cities and places from the text.

"""


# enables `response_model` in create call
client = instructor.from_openai(
    OpenAI(
        base_url="http://localhost:11434/v1",
        api_key="ollama",  # required, but unused
    ),
    mode=instructor.Mode.JSON,
)

llm_response = client.chat.completions.create(
    model="gemma2:9b",
    messages=[
        {
            "role": "user",
            "content": query,
        }
    ],
    response_model=Cities,
)

print(llm_response)

and the output is:

city='Berlin'

How should I edit my code to get the output as a list of identified cities?

[kormilitzin]

I found how to edit my code to rerun a list of cities. Thank you.

[rma-rripken]

@kormilitzin, if you are still around, can you document how you solved your problem? I think Cities needs to hold a List[str].