Files
drc-ners-nlp/processing/test.py
T

27 lines
560 B
Python

import ollama
from pydantic import BaseModel
from misc import load_prompt
class NameAnalysis(BaseModel):
identified_name: str | None
identified_surname: str | None
name = input("Enter name: ")
client = ollama.Client()
response = client.chat(
model="mistral:7b",
messages=[
{"role": "system", "content": load_prompt()},
{"role": "user", "content": name}
],
format=NameAnalysis.model_json_schema()
)
analysis = NameAnalysis.model_validate_json(response.message.content)
result = analysis.model_dump()
print(result)