File size: 1,190 Bytes
9c64cf7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
from enum import Enum
SYSTEM_PROMPT = """Tu dois classer des documents selon leur type. Les types possibles sont :
- Bon de livraison : un bon de libraison pour de la nourriture pour animaux d'élevage.
- Rapport d'analyse fourrage : un document décrivant les informations nutritionnelles et composition d'un fourrage.
Répond uniquement avec "livraison" ou "analyse". Ne réponds jamais autre chose.
"""
def classify_document (client, document_source):
model = "mistral-small-latest" # Mistral Medium 3.1
messages = [
{
"role": "system",
"content": SYSTEM_PROMPT,
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Quel est le type de ce document ?"
},
document_source
]
}
]
chat_response = client.chat.complete(
model=model,
messages=messages,
temperature=0.0,
)
if 'livraison' in chat_response.choices[0].message.content.lower():
return "livraison"
elif 'analyse' in chat_response.choices[0].message.content.lower():
return "analyse"
else:
raise ValueError("Document type could not be classified as 'livraison' or 'analyse'") |