llm_worker.py

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from multiprocessing import Queue
import torch
import time

# Limite de VRAM (opcional, apenas para log)
try:
    VRAM_TOTAL = torch.cuda.get_device_properties(0).total_memory
except Exception:
    VRAM_TOTAL = 0  # Se não houver GPU disponível

VRAM_LIMITS = {
    "llm": float(VRAM_TOTAL * 0.45),
    "stt": float(VRAM_TOTAL * 0.30),
    "tts": float(VRAM_TOTAL * 0.25)
}

def llm_worker(input_queue: Queue, output_queue: Queue):
    print("🦾 Iniciando LLM Worker...")

    model_name = "google/gemma-3-1b-it"
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Modelo começa na CPU
    model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")
    print("🧠 Modelo carregado na CPU.")

    while True:
        try:
            item = input_queue.get()
            if item is None:
                break  # Sinal de encerramento

            print("🔄 Movendo modelo para GPU...")
            model.to("cuda")

            # Cria pipeline temporária
            llm_pipeline = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=128,
                temperature=0.7,
                top_k=20,
                device=0  # Usa GPU
            )

            response = llm_pipeline(item["text"])
            generated_text = response[0]["generated_text"].replace(item["text"], "").strip()

            output_queue.put({"text": generated_text})

            print("💤 Devolvendo modelo para CPU...")
            model.to("cpu")
            torch.cuda.empty_cache()

        except Exception as e:
            print(f"🚨 Erro no processamento: {str(e)}")
            torch.cuda.empty_cache()
            time.sleep(1)