from fastapi import FastAPI, HTTPException, UploadFile, File
from transformers import pipeline
from pydantic import BaseModel
import os

app=FastAPI()

# Set the cache directory for Hugging Face transformers within /app
cache_dir = "/app/.cache"
os.environ["TRANSFORMERS_CACHE"] = cache_dir

# Load the Whisper model for transcription
transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-small")

# Pydantic model for request body validation
class TranscriptionRequest(BaseModel):
    file_name: str

# Pydantic model for response data validation
class TranscriptionResponse(BaseModel):
    text: str

# Endpoint to handle audio file transcription
@app.post("/transcribe/", response_model=TranscriptionResponse)
async def transcribe_audio(transcription_request: TranscriptionRequest):
    file_path = os.path.join("/huggingface/", "audio_files", transcription_request.file_name)

    # Check if file exists
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")

    # Read the audio file content
    with open(file_path, "rb") as audio_file:
        audio_content = audio_file.read()

    # Perform transcription using the loaded model
    transcription = transcriber(audio_content)["text"]

    return TranscriptionResponse(text=transcription)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)