keannu125/DocManagerBackend
1
0
Fork 0
mirror of https://github.com/lemeow125/DocManagerBackend.git synced 2025-01-19 01:23:02 +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" TEST_DATA = "True"
# Ollama for Categorization # Ollama for Categorization
OLLAMA_URL = "" OLLAMA_URL = "localhost:11434"
OLLAMA_USE_AUTH = "False"
OLLAMA_MODEL = "knoopx/mobile-vlm:3b-fp16"
OLLAMA_USERNAME = "" OLLAMA_USERNAME = ""
OLLAMA_PASSWORD = "" OLLAMA_PASSWORD = ""

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

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