Update README.md

README.md

@@ -9,10 +9,54 @@ base_model:
 
 
 ```python
-from transformers import pipeline
+from vllm import LLM
 
-pipe = pipeline("text-classification", model="J-LAB/Classificador_de_prompt")
-pipe("Prompt")
+# 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:
 ```json