FastAPI 如何在没有UI（即HTML）的情况下将文件（docx、doc、pdf或json）发送到FastAPI并进行预测

在本文中，我们将介绍如何在没有UI的情况下通过FastAPI发送文件，并进行预测。我们将使用FastAPI提供的文件上传功能，并结合预测模型进行文件的预测。

文件上传功能

FastAPI提供了File类来处理文件的上传。在API路由函数中，我们可以将文件参数声明为File类型。例如，以下代码演示了如何接收一个上传的文件：

from fastapi import FastAPI, File

app = FastAPI()

@app.post("/upload")
async def upload_file(file: bytes = File(...)):
    # 处理上传的文件
    return {"filename": file.filename}

在上述代码中，我们通过file: bytes = File(...)声明了一个名为file的文件参数。在这个例子中，我们仅仅返回了上传文件的文件名。

文件类型限制

在实际的应用中，我们可能只允许上传特定类型的文件，比如docx、doc、pdf或json文件。为了限制上传文件的类型，我们可以使用magic库来检查文件的类型。下面的示例代码展示了如何使用magic库来检查文件的类型。

import magic

def is_valid_file(file_path):
    blob = open(file_path, 'rb').read(2048)
    file_type = magic.from_buffer(blob, mime=True)
    valid_types = ['application/msword', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', 'application/pdf', 'application/json']
    return file_type in valid_types

通过调用magic.from_buffer(blob, mime=True)方法来获取文件的MIME类型，并检查它是否在我们允许的文件类型列表中。

文件预测

一旦我们成功地接收到文件并验证了文件类型，我们就可以使用预测模型对文件进行预测。以下示例代码展示了如何在FastAPI中使用预测模型对上传的文件进行预测。

import torch
from transformers import BertTokenizer, BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained("bert-base-uncased")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

@app.post("/predict")
async def predict(file: bytes = File(...)):
    # 处理上传的文件
    text = extract_text_from_file(file)

    # 文本预处理
    inputs = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        truncation=True,
        padding='max_length',
        max_length=128,
        return_tensors='pt'
    )

    # 模型预测
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_labels = torch.argmax(logits, dim=1)

    return {"predicted_labels": predicted_labels.tolist()}

在上述代码中，我们使用了一个预训练的BERT模型，对文本进行分类。在predict路由函数中，我们首先从文件中提取文本，然后使用BERT的tokenizer对文本进行预处理。接下来，我们使用预训练模型对处理后的文本进行预测，并返回预测结果。

运行FastAPI应用

要运行上述的FastAPI应用，我们需要在终端中执行以下命令：

uvicorn main:app --reload

其中，main是我们的主要Python文件，app是FastAPI实例的名称。--reload选项使得在代码更改时，应用能够自动重新加载。

我们可以使用curl命令来测试API，如下所示：

curl -X POST -H "Content-Type: multipart/form-data" -F "file=@example.docx" http://localhost:8000/upload

这将上传名为example.docx的文件，并返回文件名。

接下来，我们可以使用类似的方法来测试预测功能：

curl -X POST -H "Content-Type: multipart/form-data" -F "file=@example.docx" http://localhost:8000/predict

这将上传文件example.docx并返回预测结果。

总结

本文介绍了如何在没有UI（即HTML）的情况下将文件发送到FastAPI并进行预测。我们使用FastAPI提供的文件上传功能来接收文件，并使用预训练模型进行文件的预测。通过限制文件类型并使用适当的库来处理文件，我们可以确保我们只处理允许的文件类型。通过使用FastAPI和预测模型，我们可以轻松地构建一个文件处理和预测的应用程序。