import os from dotenv import load_dotenv from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from typing import List, Dict, Union, Optional # --- NOVAS IMPORTAÇÕES PARA TRANSFORMERS E PyTorch --- try: from transformers import AutoTokenizer, AutoModelForCausalLM import torch except ImportError: print("Erro: As bibliotecas 'transformers' ou 'torch' não foram encontradas.") print("Certifique-se de instalá-las: pip install transformers torch") exit(1) # Carregar variáveis de ambiente (MODEL_PATH não será mais usado diretamente aqui) load_dotenv() # --- Configurações do Modelo --- # O nome do modelo no Hugging Face Hub HF_MODEL_ID = "pierreguillou/gpt2-small-portuguese" # Instâncias globais para o modelo e tokenizer tokenizer: Optional[AutoTokenizer] = None model: Optional[AutoModelForCausalLM] = None # --- Inicialização do FastAPI --- app = FastAPI( title="AutisMind AI Server", description="API para interação com o modelo GPT2-small-portuguese e análise de conversas.", version="1.0.0", ) # --- Configuração CORS --- origins = [ "http://localhost:3000", "http://localhost:5000", "http://127.0.0.1:3000", "http://127.0.0.1:5000", # Adicione aqui os domínios do seu frontend e do Express API quando estiverem em produção # Ex: "https://seunome-seu-space-name.hf.space" ] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # --- Modelos Pydantic para Requisições e Respostas (Mantidos) --- class ChatMessage(BaseModel): role: str # "user" ou "assistant" content: str class ChatRequest(BaseModel): message: str = Field(..., description="A mensagem do usuário.") history: List[ChatMessage] = Field(default_factory=list, description="Histórico da conversa (opcional).") persona: str = Field( "Você é um assistente útil e amigável. Responda de forma clara e empática.", description="A personalidade do personagem para a IA." ) chatId: Optional[str] = Field(None, description="ID do chat para contexto (não usado pela IA, mas pode ser útil para logs).") class AnalysisResult(BaseModel): sentiment: str = Field("neutro", description="Sentimento da mensagem (positivo, negativo, neutro).") emotion: str = Field("desconhecida", description="Emoção detectada (ex: alegria, tristeza, raiva).") class ChatResponse(BaseModel): ai_response: str = Field(..., description="A resposta gerada pela IA.") analysis: AnalysisResult = Field(..., description="Resultados da análise da mensagem do usuário ou da resposta da IA.") # --- Evento de inicialização da aplicação (carrega o modelo) --- @app.on_event("startup") async def startup_event(): global model, tokenizer print(f"Carregando modelo Transformers: {HF_MODEL_ID}") try: tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID) # O GPT-2 não tem um token PAD padrão, define um. # 50256 é o token EOS padrão do GPT-2, pode ser usado como PAD também. tokenizer.pad_token = tokenizer.eos_token # Ajusta o tamanho máximo de sequência do tokenizer para algo comum como 1024 tokenizer.model_max_length = 1024 model = AutoModelForCausalLM.from_pretrained(HF_MODEL_ID) model.eval() # Coloca o modelo em modo de avaliação (desativa dropout, etc.) # Tenta mover o modelo para a GPU se disponível e compatível com PyTorch/Transformers if torch.cuda.is_available(): model.to("cuda") print("Modelo movido para GPU!") else: print("Modelo carregado na CPU.") print(f"Modelo {HF_MODEL_ID} carregado com sucesso!") except Exception as e: print(f"Erro ao carregar o modelo {HF_MODEL_ID}: {e}") raise RuntimeError(f"Falha ao carregar o modelo LLM: {e}") # --- Endpoint de Saúde --- @app.get("/health") async def health_check(): if model and tokenizer: # Verifica se ambos foram carregados return {"status": "ok", "model_loaded": True} return {"status": "loading", "model_loaded": False} # --- Endpoint Principal para Geração de Resposta e Análise --- @app.post("/generate-response", response_model=ChatResponse) async def generate_ai_response(request: ChatRequest): global model, tokenizer if not model or not tokenizer: raise HTTPException(status_code=503, detail="Modelo AI ainda não carregado. Tente novamente em breve.") user_message = request.message chat_history = request.history character_persona = request.persona # --- Construir o Prompt para GPT-2 --- # GPT-2 é um modelo generativo, não segue o formato <|system|> como Phi. # A persona e o histórico são concatenados. prompt_parts = [] # Inclui a persona no início do prompt, como uma instrução inicial para o modelo. prompt_parts.append(f"{character_persona}\n\n") # Adicionar histórico da conversa # É importante limitar o histórico para não exceder o model_max_length (1024 tokens para GPT-2 small). # Uma boa estratégia é adicionar o histórico mais recente. # O GPT-2 não tem roles, então formatamos como um diálogo. for turn in chat_history[-5:]: # Inclui os últimos 5 turnos (adaptar conforme necessidade) if turn.role == 'user': prompt_parts.append(f"Usuário: {turn.content}\n") elif turn.role == 'assistant': prompt_parts.append(f"Assistente: {turn.content}\n") # Adiciona a mensagem atual do usuário e pede para a IA responder prompt_parts.append(f"Usuário: {user_message}\nAssistente:") full_prompt = "".join(prompt_parts) try: # Codifica o prompt para tokens # return_tensors="pt" para PyTorch # max_length para garantir que o prompt não exceda o limite do modelo inputs = tokenizer( full_prompt, return_tensors="pt", max_length=tokenizer.model_max_length, truncation=True ) # Mover os inputs para a GPU se o modelo estiver na GPU if torch.cuda.is_available(): inputs = {k: v.to("cuda") for k, v in inputs.items()} # Gerar a resposta # pad_token_id: O token usado para preencher sequências, aqui usamos o EOS token do GPT-2 # do_sample=True: Habilita amostragem (para criatividade) # top_k, top_p: Métodos de amostragem para controlar a diversidade da resposta # max_new_tokens: O número máximo de novos tokens que a IA pode gerar APÓS o prompt sample_outputs = model.generate( inputs.input_ids, pad_token_id=tokenizer.pad_token_id, # Usando o token de pad que definimos do_sample=True, max_new_tokens=150, # Gera até 150 tokens novos de resposta temperature=0.7, top_k=50, top_p=0.9, num_return_sequences=1 ) # Decodificar a resposta gerada # skip_special_tokens=True: Remove tokens especiais como , , generated_text = tokenizer.decode(sample_outputs[0], skip_special_tokens=True) # --- Pós-processamento da Resposta do GPT-2 --- # GPT-2 vai gerar o prompt INTEIRO MAIS a resposta. # Precisamos remover o prompt original para obter apenas a resposta do "Assistente:". # Isso pode ser um pouco delicado e pode precisar de ajustes finos dependendo da saída real do modelo. ai_response_raw = generated_text.replace(full_prompt, "").strip() # O GPT-2 pode continuar gerando texto ou até começar um "Usuário:" de novo. # Tentamos cortar a resposta na primeira ocorrência de "Usuário:" ou "Assistente:" para evitar isso. if "\nUsuário:" in ai_response_raw: ai_response = ai_response_raw.split("\nUsuário:")[0].strip() elif "\nAssistente:" in ai_response_raw: ai_response = ai_response_raw.split("\nAssistente:")[0].strip() else: ai_response = ai_response_raw # --- Lógica de Análise (Mantida) --- analysis_result = AnalysisResult() user_message_lower = user_message.lower() if "triste" in user_message_lower or "chateado" in user_message_lower or "mal" in user_message_lower: analysis_result.sentiment = "negativo" analysis_result.emotion = "tristeza" elif "feliz" in user_message_lower or "alegre" in user_message_lower or "bom" in user_message_lower: analysis_result.sentiment = "positivo" analysis_result.emotion = "alegria" else: analysis_result.sentiment = "neutro" analysis_result.emotion = "desconhecida" return ChatResponse(ai_response=ai_response, analysis=analysis_result) except Exception as e: print(f"Erro na geração da IA: {e}") raise HTTPException(status_code=500, detail=f"Erro interno do servidor AI: {str(e)}") # --- Executar o Servidor (apenas para teste direto via Python) --- if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=5001)