如何在 Notebook 使用 FastAPI 对外提供 RestAPI，当我们测试完模型时，有事需要对外提供 API 进行测试。FastAPI 是 Python 中快速提供 Rest API 的框架，本文将对外实现一个 OCR 的图片转文字的服务。

OCR 服务

本文使用 GOT OCR 提供 OCR 识别服务，具体模型信息可以参考官网 https://github.com/Ucas-HaoranWei/GOT-OCR2.0/

安装模型依赖

!pip install tiktoken==0.6.0
!pip install verovio==4.3.1
!pip install accelerate==0.28.0

加载模型

from modelscope import AutoModel, AutoTokenizer, snapshot_download

model_id = snapshot_download("stepfun-ai/GOT-OCR2_0", local_dir="./ocr")


tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModel.from_pretrained(model_id, trust_remote_code=True, device_map="cuda", low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval().cuda()

启动 FastAPI 服务

from fastapi import FastAPI
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io
import shutil

app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello": "World"}

@app.post("/upload-image/")
async def upload_image(file: UploadFile = File(...)):
    print(file.content_type)
    # Check if the uploaded file is an image
    if file.content_type.startswith('image/'):
        with open(f"./{file.filename}", "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        extracted_text = model.chat(tokenizer, f"./{file.filename}", ocr_type='ocr')
        
        # Return the extracted text
        return {"filename": file.filename, "extracted_text": extracted_text}
    
    return {"error": "The uploaded file is not an image."}


import nest_asyncio
import uvicorn

nest_asyncio.apply()
uvicorn.run(app, port=8895, host="0.0.0.0")

进入 Swagger 页面进行测试 http://10.91.3.116:8895/docs#/default/upload_image_upload_image__post

总结

FastAPI 使用起来非常方便，可以快速的与 Notebook 代码集成，大大提高开发测试效率。