from vllm import LLM

# Inicializar o modelo
llm = LLM(
    model="J-LAB/Classificador_de_prompt", 
    task="classify", 
    enforce_eager=True, 
    hf_overrides={"architectures": ["LlamaForCausalLM"]}
)

# Lista de prompts para testar
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

# Dicionário com os labels das classes (baseado no modelo original)
id2label = {
    0: "Arte e Cultura",
    1: "Ciência e Tecnologia",
    2: "Educação e Aprendizado",
    3: "Entretenimento e Lazer",
    4: "Negócios e Finanças",
    5: "Notícias e Atualidades",
    6: "Programação e Desenvolvimento",
    7: "Saúde e Bem-estar",
    8: "Sociedade e Relacionamentos",
    9: "Viagens e Turismo"
}

# Classificar os prompts
outputs = llm.classify(prompts)

# Imprimir resultados com os labels
for prompt, output in zip(prompts, outputs):
    probs = output.outputs.probs
    
    # Criar lista de tuplas (probabilidade, classe) e ordenar
    prob_class_pairs = [(prob, id2label[i]) for i, prob in enumerate(probs)]
    prob_class_pairs.sort(reverse=True)
    
    # Mostrar as 3 classes mais prováveis
    print(f"\nPrompt: {prompt!r}")
    print("Top 3 classes mais prováveis:")
    for prob, class_name in prob_class_pairs[:3]:
        print(f"- {class_name}: {prob:.2%}")

Categorias:

    "0": "Analise e Interpretação",
    "1": "Aprendizado e Tutorial",
    "2": "Consultoria e Estratégia",
    "3": "Criação e Criatividade",
    "4": "Cálculos e Matemática",
    "5": "Informação Geral",
    "6": "Programação e Desenvolvimento",
    "7": "Suporte e Ajuda"

Accuracy: 88,80%

Downloads last month
110
Safetensors
Model size
435M params
Tensor type
F32
·
Inference Providers NEW
This model is not currently available via any of the supported third-party Inference Providers, and the model is not deployed on the HF Inference API.

Model tree for J-LAB/Classificador_de_prompt

Finetuned
(11)
this model