Beega13/SillyTavernBot

1import gradio as gr2from transformers import AutoModelForCausalLM, AutoTokenizer3import torch4 5# Modelo escolhido para roleplay6MODEL_NAME = "PygmalionAI/pygmalion-6b"7 8# Carregando tokenizer e modelo9tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)10model = AutoModelForCausalLM.from_pretrained(11    MODEL_NAME,12    torch_dtype=torch.float16,13    device_map="auto"14)15 16# Função de chat17def chat(message, history):18    # Juntar histórico em uma string para dar mais contexto19    history_text = ""20    for user_msg, bot_msg in history:21        history_text += f"User: {user_msg}\nBot: {bot_msg}\n"22    prompt = history_text + f"User: {message}\nBot:"23 24    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)25    outputs = model.generate(26        inputs,27        max_length=300,28        pad_token_id=tokenizer.eos_token_id,29        do_sample=True,30        temperature=0.8,31        top_p=0.932    )33    response = tokenizer.decode(outputs[0], skip_special_tokens=True)34 35    # Pegar apenas a parte da resposta depois de "Bot:"36    if "Bot:" in response:37        response = response.split("Bot:")[-1].strip()38 39    return response40 41# Interface Gradio42iface = gr.ChatInterface(43    fn=chat,44    title="SillyTavernBot (sem censura)",45    description="Chat experimental rodando com Pygmalion-6B"46)47 48if __name__ == "__main__":49    iface.launch(server_name="0.0.0.0", server_port=7860)50