File size: 1,762 Bytes
c3b1078 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from pathlib import Path
import os
from gector import GecBERTModel
from faster_whisper import WhisperModel, BatchedInferencePipeline
from transformers.models.whisper.english_normalizer import BasicTextNormalizer
from text_processing.inverse_normalize import InverseNormalizer
import shutil
import uvicorn
# Initialize the FastAPI app
app = FastAPI()
# Initialize models and normalizer
current_dir = Path(__file__).parent.as_posix()
inverse_normalizer = InverseNormalizer('vi')
whisper_model = WhisperModel("pho_distill_q8", device="auto", compute_type="auto")
batched_model = BatchedInferencePipeline(model=whisper_model, use_vad_model=True, chunk_length=15)
gector_model = GecBERTModel(
vocab_path=os.path.join(current_dir, "gector/vocabulary"),
model_paths=[os.path.join(current_dir, "gector/Model_GECTOR")],
split_chunk=True
)
normalizer = BasicTextNormalizer()
@app.post("/transcriptions")
async def transcribe_audio(file: UploadFile = File(...)):
# Save the uploaded file temporarily
temp_file_path = Path(f"temp_{file.filename}")
with open(temp_file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
segments, info = batched_model.transcribe(str(temp_file_path), language="vi", batch_size=32)
os.remove(temp_file_path)
transcriptions = [segment.text for segment in segments]
normalized_transcriptions = [inverse_normalizer.inverse_normalize(normalizer(text)) for text in transcriptions]
corrected_texts = gector_model(normalized_transcriptions)
return JSONResponse({"text": ' '.join(corrected_texts)})
if __name__ == "__main__":
uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True) |