llm.py

# llm.py
import os
import base64
from PIL import Image
from io import BytesIO
import torch
from transformers import AutoProcessor, Llama4ForConditionalGeneration

READER_MODEL = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
CODER_MODEL = "tiiuae/falcon-7b-instruct"

READER_SYS_PROMPT = (
    "Você é um assistente virtual especializado em interpretação de fluxogramas. "
    "Você receberá uma imagem de um fluxograma e deve interpretá-lo, respondendo com um pseudocódigo em português."
)

CODER_SYS_PROMPT = """
Você é um assistente prestativo que converte pseudocódigos em código MicroPython para a placa BitDogLab.
Inclua sempre as bibliotecas necessárias, como utime, neopixel, SoftI2C, ssd1306, etc.
"""

class Reader:
    def __init__(self):
        self.processor = AutoProcessor.from_pretrained(READER_MODEL)
        self.model = Llama4ForConditionalGeneration.from_pretrained(
            READER_MODEL,
            attn_implementation="flex_attention",
            device_map="auto",
            torch_dtype=torch.bfloat16,
        )

    def read(self, question: str, image: Image.Image) -> str:
        img_copy = image.convert("RGB")
        img_copy.thumbnail((400, 300), Image.Resampling.LANCZOS)

        buffered = BytesIO()
        img_copy.save(buffered, format="JPEG", quality=50)
        img_bytes = buffered.getvalue()

        # Prepara mensagens no formato do LLaMA-4
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": img_bytes},
                    {"type": "text", "text": question or "Gere um pseudocódigo a partir do fluxograma da imagem."},
                ],
            },
        ]

        inputs = self.processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(self.model.device)

        outputs = self.model.generate(
            **inputs,
            max_new_tokens=512,
        )

        generated_text = self.processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]

        return f"{READER_SYS_PROMPT}\n\nUsuário: {question}\n\nResposta: {generated_text}"


class Coder:
    def __init__(self):
        self.api_url = f"https://api-inference.huggingface.co/models/{CODER_MODEL}"
        self.headers = {
            "Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}",
            "Content-Type": "application/json"
        }

    def generate(self, pseudocode: str) -> str:
        prompt = f"{CODER_SYS_PROMPT}\n\nUsuário: {pseudocode}\n\nCódigo:"
        payload = {
            "inputs": prompt,
            "parameters": {"max_new_tokens": 512}
        }

        import requests
        response = requests.post(self.api_url, headers=self.headers, json=payload)

        if response.status_code != 200:
            return f"❌ Erro ao acessar o modelo: {response.status_code} - {response.text}"

        try:
            return response.json()[0]["generated_text"]
        except (KeyError, IndexError, TypeError):
            return "⚠️ O modelo ainda está carregando ou houve um erro ao gerar o código. Tente novamente em instantes."