Desplegar Llama 2 en Google Colab y acceder mediante FAST-API y el tunnel Ngrok

Tiempo de lectura: 4 minutos

Hola, hoy vamos a aprender cómo podemos desplegar gratis el modelo de IA Llama 2 usando Google Colab y Fast API. Me baso en el tutorial que encontré navegando https://medium.com/@yuhongsun96/host-a-llama-2-api-on-gpu-for-free-a5311463c183

En este ejemplo crearé un pequeño programa que se desplegará en Google Colab.

Lo primero que vamos a hacer es acceder a Google Colab: https://colab.research.google.com/?hl=es

Y creamos un nuevo cuaderno:

Ahora le pondremos un nombre y empezaremos a añadir nuestro código:

Primero se instalan las dependencias necesarias:

llama-cpp-python y además activamos la opción de GPU en la compilación:

!CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python

Ahora instalamos fastapi y los paquetes que vamos a utilizar para crear el api REST:

!pip install fastapi[all] uvicorn python-multipart transformers pydantic tensorflow

Finalmente, se descarga el paquete de Ngrok que nos permitirá crear un tunnel para realizar conexiones a la API.

!wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip
!unzip -o ngrok-stable-linux-amd64.zip

Ahora tenemos que registrarnos en Ngrok para que podamos acceder a la API URL de forma pública:

https://dashboard.ngrok.com/signup

Creamos una cuenta gratuita, que permite obtener 1 tunnel y obtenemos el token de acceso:

Y ahora copiamos el token y lo utilizamos añadiendo esta línea de código:

!./ngrok authtoken <TOKEN>

Y ahora vamos a implementar toda la lógica de API rest con FAST-API:

%%writefile app.py
from typing import Any

from fastapi import FastAPI
from fastapi import HTTPException
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import tensorflow as tf


# Descargar el modelo de la web higginface
HUGGINGFACE_REPO= "TheBloke/Llama-2-13B-chat-GGUF"
HUGGINGFACE_AI_MODEL= "llama-2-13b-chat.Q5_0.gguf"

model_path = hf_hub_download(
    repo_id=HUGGINGFACE_REPO,
    filename=HUGGINGFACE_AI_MODEL
)

llama2_model = Llama(
    model_path=model_path,
    n_gpu_layers=64,
    n_ctx=2000
)

app = FastAPI()


class TextInput(BaseModel):
    inputs: str
    parameters: dict[str, Any] | None


@app.get("/")
def status_gpu_check() -> dict[str, str]:
    gpu_msg = "Available" if tf.test.is_gpu_available() else "Unavailable"
    return {
        "status": "I am ALIVE!",
        "gpu": gpu_msg
    }


@app.post("/enviar_pregunta/")
async def generate_text(data: TextInput) -> dict[str, str]:
    try:
        params = data.parameters or {}
        response = llama2_model(prompt=data.inputs, **params)
        model_out = response['choices'][0]['text']
        return {"generated_text": model_out}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

*Recuerda que los modelos que puedes utilizar con llama 2 ahora son en formato gguf

Ahora ejecutamos el servidor:

!uvicorn app:app --host 0.0.0.0 --port 8000 > server.log 2>&1 &

Podemos ver los logs:

!curl localhost:8000

Ahora enlazamos Ngrok

from IPython import get_ipython
get_ipython().system_raw('./ngrok http 8000 &')

Y ahora vamos a comprobar si se ha generado correctamente el tunnel:

!curl -s http://localhost:4040/api/tunnels | python3 -c "import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])"

Cuando queramos parar los servicios tendremos que ejecutar este código:

!pkill uvicorn
!pkill ngrok

Ahora vamos a utilizar Google T4.

Vamos al menú de la derecha y seleccionamos la flecha:

Elegimos cambiar entorno de ejecución:

Y elegimos T4 GPU, pulsamos en guardar:

Y pulsamos en conectar:

Ahora, una vez conectado, podemos ejecutar todo el entorno:

Ahora vamos ejecutando cada línea de código:

Miraremos el log, para esperar a que termine de descargar el modelo nada más arrancamos el server.

Una vez ejecutado, podremos acceder a la API desde esta URL:

Cuando esté listo ejecutamos: !curl localhost:8000 y comprobaremos que devuelve el mensaje:

Y ahora para obtener la URL ejecutamos este código:

!curl -s http://localhost:4040/api/tunnels | python3 -c "import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])"

Y devolverá la URL que nos permitirá acceder.

Ahora podremos acceder a la URL de endpoints con esa dirección, y si añadimos /docs accederemos a la interfaz de swagger.

Le damos a ejecutar y podremos obtener la respuesta:

*Nota: Google Colab en versión gratis, solo permite ejecución de 12 horas seguidas. Esta solución no permite hospedar esta API para un uso continuado. Siempre podremos adquirir la versión de pago de Google Colab.