main.py

# ------------------------------------ IMPORTS ------------------------------------

from fastapi import FastAPI

from fastapi.middleware.cors import CORSMiddleware
import os

from pytube import YouTube #Descarga de audio

import whisper #Transcripcion

from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer #Traduccion

import torch #GPU

# ------------------------------------ VARIABLES ------------------------------------

nombre_archivo = "audio.mp3" # Nombre del archivo de audio
app = FastAPI() # Inicializacion de la API

origins = [
    "http://localhost",
    "http://localhost:8080",
    "http://localhost:3000",
    "http://localhost:8000",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=['*'],
    allow_headers=["*"],
)

# Revisa si hay GPU
if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")


# ------------------------------------ FUNCIONES ------------------------------------
# Post
@app.get("/")
def index():
    return "hola"

@app.post("/descargar/{url}")
def descargar_audio(url):
    url = "https://www.youtube.com/watch?v=" + url
    yt = YouTube(url)
    ruta_archivo = "/data" + nombre_archivo  # Ruta específica donde deseas guardar el archivo
    yt.streams.filter(only_audio=True).first().download(output_path=ruta_archivo, filename=nombre_archivo)
    return {"url": "Listo"}

@app.get("/transcribir") 
def transcribir():
    print("Transcribiendo...")
    model = whisper.load_model("medium", device=device)
    result = model.transcribe(nombre_archivo, fp16=False)
    transcipcion = result["text"]
    print(transcipcion)
    return {"transcripcion": transcipcion}

@app.get("/traducir/{idioma}/{idioma_destino}")
def traducir(idioma,idioma_destino):

    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_1.2B").to(device)
    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_1.2B")

    f = open ('texto.txt','r')
    texto = f.read()
    f.close()

    listaPalabras = texto.split(" ")
    lista = [] # Lista para almacenar las palabras que se van a traducir
    palabras_sobrantes = []  # Lista para almacenar las palabras restantes
    result_stringF = ""

    i = 0
    while i < len(listaPalabras): # Recorre la lista de palabras para separar en grupos de 100

        if i % 100 == 0:
            textoResult = " ".join(lista)
            tokenizer.src_lang = idioma
            encoded_hi = tokenizer(textoResult, return_tensors="pt").to(device)
            generated_tokens = model.generate(**encoded_hi, forced_bos_token_id=tokenizer.get_lang_id(idioma_destino))
            #Añade resultados a un string
            results = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
            result_string = " ".join(results)
            result_stringF = result_stringF + result_string
            lista = []

        lista.append(listaPalabras[i])

        if i == len(listaPalabras) - 1:
            palabras_sobrantes.extend(lista)  # Almacena las palabras restantes al final del bucle

        i = i + 1

    if len(palabras_sobrantes) > 0:
        textoResult = " ".join(palabras_sobrantes)
        tokenizer.src_lang = idioma
        encoded_hi = tokenizer(textoResult, return_tensors="pt").to(device)
        generated_tokens = model.generate(**encoded_hi, forced_bos_token_id=tokenizer.get_lang_id(idioma_destino))
        #Añade resultados a un string
        results = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
        result_string = " ".join(results)
        result_stringF = result_stringF + result_string


    return {"Traduccion": results}