Update app.py

app.py

@@ -1,63 +1,109 @@
 
 import os
-import json
 import numpy as np
 import gradio as gr
-
-# ---------- 1) Sentiment Baseline (TF-IDF + LogisticRegression) ----------
 import joblib
+
 from sklearn.pipeline import Pipeline
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
 
+from transformers import pipeline as hf_pipeline
+
+
+# ======================================================================
+# 1. Baseline de Sentimentos (TF-IDF + Logistic Regression)
+# ======================================================================
+
 BASELINE_PATH = os.getenv("MODEL_PATH", "baseline_pipe.pkl")
 
-def train_small_baseline(save_path=BASELINE_PATH, sample_pct="train[:0.5%]"):
+
+def train_small_baseline(save_path: str = BASELINE_PATH,
+                         sample_split: str = "train[:0.5%]"):
+    """
+    Treina um baseline pequeno usando uma amostra do dataset amazon_polarity.
+    Usado apenas se baseline_pipe.pkl nao existir no Space.
+    """
     from datasets import load_dataset
     import pandas as pd
-    ds = load_dataset("amazon_polarity", split=sample_pct)
+
+    ds = load_dataset("amazon_polarity", split=sample_split)
     df = pd.DataFrame({"text": ds["content"], "label": ds["label"]})
-    pipe = Pipeline([
-        ("tfidf", TfidfVectorizer(max_features=30000, ngram_range=(1,2))),
-        ("clf", LogisticRegression(max_iter=1000)),
-    ])
+
+    pipe = Pipeline(
+        [
+            ("tfidf", TfidfVectorizer(max_features=30000, ngram_range=(1, 2))),
+            ("clf", LogisticRegression(max_iter=1000)),
+        ]
+    )
+
     pipe.fit(df["text"], df["label"])
     joblib.dump(pipe, save_path)
     return pipe
 
+
 def load_or_bootstrap_baseline():
+    """
+    Se existir baseline_pipe.pkl, carrega.
+    Se nao existir e DISABLE_AUTOTRAIN != 1, treina um baseline pequeno.
+    """
     if os.path.exists(BASELINE_PATH):
         return joblib.load(BASELINE_PATH)
-    if os.getenv("DISABLE_AUTOTRAIN", "0") == "1":
+
+    disable_auto = os.getenv("DISABLE_AUTOTRAIN", "0")
+    if disable_auto == "1":
         return None
+
     return train_small_baseline()
 
-baseline = load_or_bootstrap_baseline()
 
-def classify_only(text):
-    if not text or not text.strip():
+baseline_model = load_or_bootstrap_baseline()
+
+
+def classify_only(text: str):
+    """
+    Apenas classifica o sentimento (positivo/negativo) e retorna JSON.
+    """
+    if not text or text.strip() == "":
         return {"erro": "Digite um texto."}
-    if baseline is None:
-        return {"erro": "Modelo baseline nao encontrado. Envie baseline_pipe.pkl ou remova DISABLE_AUTOTRAIN."}
-    proba = baseline.predict_proba([text])[0]
+
+    if baseline_model is None:
+        return {
+            "erro": (
+                "Modelo baseline nao encontrado. "
+                "Envie baseline_pipe.pkl na aba Files ou remova DISABLE_AUTOTRAIN."
+            )
+        }
+
+    proba = baseline_model.predict_proba([text])[0]
     pred = int(np.argmax(proba))
     label = "positivo" if pred == 1 else "negativo"
     conf = float(np.max(proba))
+
     return {"sentimento": label, "confianca": round(conf, 3)}
 
-# ---------- 2) Generative Assistant (FLAN-T5) ----------
-from transformers import pipeline
+
+# ======================================================================
+# 2. IA Generativa (FLAN-T5) para resposta ao cliente
+# ======================================================================
+
 GEN_MODEL_ID = os.getenv("GEN_MODEL_ID", "google/flan-t5-base")
-generator = pipeline("text2text-generation", model=GEN_MODEL_ID)
+
+generator = hf_pipeline("text2text-generation", model=GEN_MODEL_ID)
 
 SYSTEM_PROMPT = (
     "Voce e um atendente virtual educado que responde em portugues do Brasil. "
     "Use poucas frases, tom empatico e pratico."
 )
 
-def generate_reply(user_text, sentimento_json):
-    if not user_text or not user_text.strip():
+
+def generate_reply(user_text: str, sentimento_json):
+    """
+    Gera uma resposta em PT-BR condicionada ao sentimento detectado.
+    """
+    if not user_text or user_text.strip() == "":
         return "Digite uma mensagem."
+
     sentimento = None
     if isinstance(sentimento_json, dict) and "sentimento" in sentimento_json:
         sentimento = sentimento_json["sentimento"]
@@ -65,53 +111,91 @@ def generate_reply(user_text, sentimento_json):
     if sentimento == "negativo":
         intent = (
             "A avaliacao do cliente e NEGATIVA. "
-            "Peca desculpas, reconheca o problema, ofereca ajuda objetiva e solicite informacoes adicionais."
+            "Peca desculpas, reconheca o problema, ofereca ajuda objetiva "
+            "e solicite informacoes adicionais (numero do pedido, produto, contato)."
         )
     elif sentimento == "positivo":
         intent = (
             "A avaliacao do cliente e POSITIVA. "
-            "Agradeca de forma calorosa, reforce pontos fortes mencionados e convide para novas compras."
+            "Agradeca de forma calorosa, reforce os pontos positivos citados "
+            "e convide o cliente a continuar comprando."
         )
     else:
-        intent = "O sentimento nao foi identificado. Responda de forma neutra, util e cordial."
+        intent = (
+            "O sentimento nao foi identificado. "
+            "Responda de forma neutra, cordial e util."
+        )
 
     prompt = (
-        f"{SYSTEM_PROMPT}\n\n{intent}\n\n"
+        f"{SYSTEM_PROMPT}\n\n"
+        f"{intent}\n\n"
         f"Mensagem do cliente:\n\"{user_text}\"\n\n"
-        f"Escreva uma resposta curta (2-4 frases)."
+        f"Escreva uma resposta curta (2 a 4 frases)."
     )
-    out = generator(prompt, max_length=128, do_sample=True, temperature=0.6, top_p=0.9)[0]["generated_text"]
+
+    out = generator(
+        prompt,
+        max_length=128,
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.9,
+    )[0]["generated_text"]
+
     return out
 
-# ---------- 3) Gradio UI ----------
-with gr.Blocks(title="Classificador + Chatbot (Sentimentos)") as demo:
+
+# ======================================================================
+# 3. Interface Gradio - duas abas (Analise e Chatbot)
+# ======================================================================
+
+with gr.Blocks(title="Chatbot de Sentimentos - Professor Rodrigo") as demo:
     gr.Markdown(
-        \"\"\"
-        # Chatbot de Sentimentos (ML + IA Generativa)
-        Professor Rodrigo - Projeto Final ML & DL
+        """
+# Chatbot de Sentimentos (ML + IA Generativa)
+**Professor Rodrigo** - Projeto Final ML & DL
 
-        - Classificacao: TF-IDF + Regressao Logistica (autotreino minimo se nao houver `baseline_pipe.pkl`).  
-        - Geracao: `google/flan-t5-base` para redigir respostas educadas em PT-BR.
+- Classificacao: TF-IDF + Regressao Logistica (baseline).
+- Geracao: `google/flan-t5-base` para redigir respostas em PT-BR.
 
-        > Dica: Envie `baseline_pipe.pkl` em *Files* para pular o autotrain.
-        \"\"\"
+> Dica: envie `baseline_pipe.pkl` na aba *Files* para usar um modelo treinado por voce.
+        """
     )
 
+    # Aba 1 - somente analise de sentimento
     with gr.Tab("Analise de Sentimento"):
-        in_text = gr.Textbox(label="Digite uma avaliacao de produto", lines=4, placeholder="Ex.: O produto chegou rapido e funciona muito bem.")
-        out_json = gr.JSON(label="Resultado")
-        btn = gr.Button("Analisar")
-        btn.click(classify_only, inputs=in_text, outputs=out_json)
+        input_text = gr.Textbox(
+            label="Digite uma avaliacao de produto",
+            lines=4,
+            placeholder="Ex.: O produto chegou rapido e superou minhas expectativas.",
+        )
+        output_json = gr.JSON(label="Resultado da classificacao")
+        btn_analisar = gr.Button("Analisar")
+        btn_analisar.click(classify_only, inputs=input_text, outputs=output_json)
 
+    # Aba 2 - Chatbot (analise + resposta generativa)
     with gr.Tab("Chatbot (Analise + Resposta)"):
-        chat_in = gr.Textbox(label="Mensagem do cliente", lines=4)
+        chat_input = gr.Textbox(
+            label="Mensagem do cliente",
+            lines=4,
+            placeholder="Ex.: Estou chateado, o produto chegou quebrado.",
+        )
         with gr.Row():
-            analyze_btn = gr.Button("1) Analisar sentimento")
-            gen_btn = gr.Button("2) Gerar resposta")
-        analysis_box = gr.JSON(label="Sentimento detectado")
-        reply_box = gr.Textbox(label="Resposta gerada", lines=6)
-        analyze_btn.click(classify_only, inputs=chat_in, outputs=analysis_box)
-        gen_btn.click(generate_reply, inputs=[chat_in, analysis_box], outputs=reply_box)
+            btn_analise_chat = gr.Button("1) Analisar sentimento")
+            btn_resposta = gr.Button("2) Gerar resposta")
+
+        chat_analysis = gr.JSON(label="Sentimento detectado")
+        chat_reply = gr.Textbox(
+            label="Resposta gerada pela IA",
+            lines=6,
+        )
+
+        btn_analise_chat.click(
+            classify_only, inputs=chat_input, outputs=chat_analysis
+        )
+        btn_resposta.click(
+            generate_reply, inputs=[chat_input, chat_analysis], outputs=chat_reply
+        )
+
 
 if __name__ == "__main__":
     demo.launch()