Nvidia LLM

from fastapi import FastAPI, File, UploadFile
import shutil, os
from ocr_processor import ocr_from_pdf
from llama_cpp import Llama

app = FastAPI()
MODEL_PATH = "./models/Mistral-Nemo-Instruct-2407-Q4_K_M.gguf"
llm = Llama(model_path=MODEL_PATH, n_ctx=8192, n_gpu_layers=41, verbose=True)

@app.post("/upload/")
async def upload_pdf(file: UploadFile = File(...)):
    temp_path = f"/tmp/{file.filename}"
    with open(temp_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    ocr_text = ocr_from_pdf(temp_path)
    os.remove(temp_path)

    result = llm(ocr_text)
    return {"output": result["choices"][0]["text"].strip()}
from fastapi import FastAPI, File, UploadFile
import shutil, os
from ocr_processor import ocr_from_pdf
from llama_cpp import Llama

app = FastAPI()
MODEL_PATH = "./models/Mistral-Nemo-Instruct-2407-Q4_K_M.gguf"
llm = Llama(model_path=MODEL_PATH, n_ctx=8192, n_gpu_layers=41, verbose=True)

@app.post("/upload/")
async def upload_pdf(file: UploadFile = File(...)):
    temp_path = f"/tmp/{file.filename}"
    with open(temp_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    ocr_text = ocr_from_pdf(temp_path)
    os.remove(temp_path)

    result = llm(ocr_text)
    return {"output": result["choices"][0]["text"].strip()}
FROM nvidia/cuda:12.3.1-devel-ubuntu22.04
SHELL ["/bin/bash", "-c"]

RUN apt-get update && apt-get install -y \
    python3-dev python3-pip \
    curl \
    build-essential \
    software-properties-common \
    tesseract-ocr \
    poppler-utils \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    && rm -rf /var/lib/apt/lists/*

RUN apt-get update && apt-get install -y ninja-build git

ENV CUDACXX=/usr/local/cuda-12/bin/nvcc
ENV CMAKE_ARGS="-DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=all-major"

WORKDIR /app
COPY . /app

RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir -r requirements.txt
services:
  llm-ocr:
    build: .
    container_name: llm_ocr_gpu
    volumes:
      - ./models:/app/models
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['0']
              capabilities: [gpu]
    ports:
      - "8000:8000"
    stdin_open: true
    tty: true
    command: ["uvicorn", "api_server:app", "--host", "0.0.0.0", "--port", "8000"]
from llama_cpp import Llama

print('Loading model...')
MODEL_PATH = "./models/Mistral-Nemo-Instruct-2407-Q4_K_M.gguf"
llm = Llama(model_path=MODEL_PATH, n_ctx=8192, n_gpu_layers=41, verbose=True)
print('Model loaded.')
from fastapi import FastAPI, File, UploadFile, Header, HTTPException
import shutil, os
from ocr_processor import ocr_from_pdf
from llama_cpp import Llama

app = FastAPI()
MODEL_PATH = "./models/Mistral-Nemo-Instruct-2407-Q4_K_M.gguf"
llm = Llama(model_path=MODEL_PATH, n_ctx=8192, n_gpu_layers=41, verbose=True)

API_TOKEN = "supersecret"  # Change this to a secure token

def authorize(authorization: str):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid authorization header")
    token = authorization.split(" ")[1]
    if token != API_TOKEN:
        raise HTTPException(status_code=403, detail="Invalid token")

@app.post("/upload/")
async def upload_pdf(file: UploadFile = File(...), authorization: str = Header(None)):
    authorize(authorization)

    temp_path = f"/tmp/{file.filename}"
    with open(temp_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    ocr_text = ocr_from_pdf(temp_path)
    os.remove(temp_path)

    result = llm(ocr_text)
    return {"output": result["choices"][0]["text"].strip()}
from fastapi import Request, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware

WHITELIST = {"127.0.0.1", "192.168.1.100"}  # Set your safe IPs
BLACKLIST = {"10.0.0.5", "172.16.0.8"}      # Block these IPs

class IPFilterMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        client_ip = request.client.host
        if client_ip in BLACKLIST:
            raise HTTPException(status_code=403, detail="IP blocked")
        if WHITELIST and client_ip not in WHITELIST:
            raise HTTPException(status_code=403, detail="IP not in whitelist")
        return await call_next(request)
from fastapi import FastAPI, Depends, File, UploadFile, HTTPException, Security
from fastapi.security import OAuth2PasswordBearer, SecurityScopes
from starlette.status import HTTP_403_FORBIDDEN
import shutil, os
from ocr_processor import ocr_from_pdf
from llama_cpp import Llama
from ip_filter import IPFilterMiddleware

app = FastAPI()
app.add_middleware(IPFilterMiddleware)

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token", scopes={"read": "Read OCR", "write": "Run OCR + LLM"})

FAKE_TOKENS = {
    "admin-token": {"scopes": ["read", "write"]},
    "read-only-token": {"scopes": ["read"]}
}

MODEL_PATH = "./models/Mistral-Nemo-Instruct-2407-Q4_K_M.gguf"
llm = Llama(model_path=MODEL_PATH, n_ctx=8192, n_gpu_layers=41, verbose=True)

def get_current_user(security_scopes: SecurityScopes, token: str = Depends(oauth2_scheme)):
    if token not in FAKE_TOKENS:
        raise HTTPException(status_code=401, detail="Invalid token")
    token_scopes = FAKE_TOKENS[token]["scopes"]
    for scope in security_scopes.scopes:
        if scope not in token_scopes:
            raise HTTPException(status_code=403, detail="Not enough permissions")
    return token

@app.post("/upload/")
async def upload_pdf(file: UploadFile = File(...), token: str = Security(get_current_user, scopes=["write"])):
    temp_path = f"/tmp/{file.filename}"
    with open(temp_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    ocr_text = ocr_from_pdf(temp_path)
    os.remove(temp_path)

    result = llm(ocr_text)
    return {"output": result["choices"][0]["text"].strip()}

/////////
For server_api.py testing
from fastapi.security import OAuth2PasswordRequestForm

@app.post("/token")
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
    if form_data.username == "admin" and form_data.password == "secret":
        return {"access_token": "admin-token", "token_type": "bearer"}
    elif form_data.username == "reader" and form_data.password == "read":
        return {"access_token": "read-only-token", "token_type": "bearer"}
    else:
        raise HTTPException(status_code=400, detail="Invalid credentials")