import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline import warnings warnings.filterwarnings("ignore") class MiMoChatBot: def __init__(self): self.model_name = "XiaomiMiMo/MiMo-V2-Flash" self.tokenizer = None self.model = None self.pipeline = None self.device = "cuda" if torch.cuda.is_available() else "cpu" def load_model(self): """Carrega o modelo e tokenizer""" try: print(f"Carregando modelo {self.model_name}...") self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.model = AutoModelForCausalLM.from_pretrained( self.model_name, torch_dtype=torch.float16 if self.device == "cuda" else torch.float32, device_map="auto" if self.device == "cuda" else None, trust_remote_code=True ) if self.device == "cpu": self.model = self.model.to(self.device) # Configura o pipeline self.pipeline = pipeline( "text-generation", model=self.model, tokenizer=self.tokenizer, max_new_tokens=512, temperature=0.7, top_p=0.95, do_sample=True, pad_token_id=self.tokenizer.eos_token_id ) print("Modelo carregado com sucesso!") return True except Exception as e: print(f"Erro ao carregar o modelo: {e}") return False def generate_response(self, message, history): """Gera uma resposta baseada na mensagem e histórico""" if not self.pipeline: return "Modelo ainda está carregando. Por favor, aguarde..." try: # Constrói o contexto com histórico context = "" for user_msg, bot_msg in history[-3:]: # Usa últimas 3 interações context += f"Usuário: {user_msg}\nAssistente: {bot_msg}\n" context += f"Usuário: {message}\nAssistente: " # Gera a resposta response = self.pipeline( context, max_new_tokens=512, temperature=0.7, do_sample=True, pad_token_id=self.tokenizer.eos_token_id, eos_token_id=self.tokenizer.eos_token_id, ) # Extrai apenas a nova parte da resposta generated_text = response[0]['generated_text'] new_response = generated_text[len(context):].strip() # Limpa a resposta se necessário if new_response.startswith("Assistente: "): new_response = new_response[len("Assistente: "):] return new_response if new_response else "Desculpe, não consegui gerar uma resposta." except Exception as e: return f"Erro ao gerar resposta: {str(e)}" # Inicializa o chatbot chatbot = MiMoChatBot() def create_chatbot_app(): """Cria a interface do aplicativo Gradio""" # Função principal do chat def chat_function(message, history): if not message.strip(): return history, "" # Adiciona a mensagem do usuário ao histórico history.append({"role": "user", "content": message}) # Gera a resposta response = chatbot.generate_response(message, [{"role": h["role"], "content": h["content"]} for h in history[:-1]]) # Adiciona a resposta do assistente history.append({"role": "assistant", "content": response}) return history, "" # Função para limpar o chat def clear_chat(): return [], "" # Carrega o modelo em background model_status = gr.Textbox("Carregando modelo...", visible=False) # 🚨 Gradio 6: NO parameters in gr.Blocks() constructor! with gr.Blocks() as demo: # Header com branding gr.HTML("""
Converse com o modelo de linguagem XiaomiMiMo/MiMo-V2-Flash