keannu125/DocManagerBackend
1
0
Fork 0
mirror of https://github.com/lemeow125/DocManagerBackend.git synced 2025-01-18 17:13:00 +08:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update Ollama API to use different model

Browse source
This commit is contained in:
Keannu Bernasol 2024-12-07 14:03:17 +08:00
parent f39f5966d6
commit d81319c8ec
2 changed files with 8 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.env.sample
View file

@ -23,6 +23,8 @@ ADMIN_PASSWORD = ''
TEST_DATA = "True"
# Ollama for Categorization
OLLAMA_URL = ""
OLLAMA_URL = "localhost:11434"
OLLAMA_USE_AUTH = "False"
OLLAMA_MODEL = "knoopx/mobile-vlm:3b-fp16"
OLLAMA_USERNAME = ""
OLLAMA_PASSWORD = ""

10
docmanager_backend/config/management/commands/start_watcher.py
View file

@ -1,7 +1,7 @@
from ollama import ChatResponse
import base64
import httpx
from django.core.management.base import BaseCommand, CommandError
from django.core.management.base import BaseCommand
from io import BytesIO
from documents.models import Document
@ -89,7 +89,7 @@ class PDFHandler(FileSystemEventHandler):
client = Client(
host=get_secret("OLLAMA_URL"),
auth=httpx.BasicAuth(
username=get_secret("OLLAMA_USERNAME"), password=get_secret("OLLAMA_PASSWORD"))
username=get_secret("OLLAMA_USERNAME"), password=get_secret("OLLAMA_PASSWORD")) if get_secret("OLLAMA_USE_AUTH") else None
)
encoded_image = base64.b64encode(
@ -111,15 +111,15 @@ class PDFHandler(FileSystemEventHandler):
"""
response: ChatResponse = client.chat(
model="llama3.2-vision",
model=get_secret("OLLAMA_MODEL"),
messages=[
{"role": "user", "content": content,
"images": [encoded_image]},
],
)
document_type = response["message"]["content"].split(":")[
0].replace("*", "").replace(".", "")
document_type = response["message"]["content"].replace(
"*", "").replace(".", "")
# A few safety checks if the model does not follow through with output instructions
if len(document_type) > 16: