Spaces:
Runtime error
Runtime error
Update main.py
Browse files
main.py
CHANGED
@@ -0,0 +1,139 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Instalação das dependências necessárias
|
2 |
+
!pip install transformers torch gradio datasets
|
3 |
+
|
4 |
+
import torch
|
5 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments
|
6 |
+
import gradio as gr
|
7 |
+
import json
|
8 |
+
import os
|
9 |
+
|
10 |
+
class GameAIAssistant:
|
11 |
+
def __init__(self, model_name="deepseek-ai/DeepSeek-R1"):
|
12 |
+
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
13 |
+
self.model = AutoModelForCausalLM.from_pretrained(model_name)
|
14 |
+
self.memory = []
|
15 |
+
self.max_memory = 10
|
16 |
+
self.game_knowledge_file = "game_knowledge.txt"
|
17 |
+
|
18 |
+
# Carregar conhecimento existente do jogo
|
19 |
+
self.load_game_knowledge()
|
20 |
+
|
21 |
+
def load_game_knowledge(self):
|
22 |
+
if os.path.exists(self.game_knowledge_file):
|
23 |
+
with open(self.game_knowledge_file, 'r', encoding='utf-8') as f:
|
24 |
+
self.game_knowledge = f.read()
|
25 |
+
else:
|
26 |
+
self.game_knowledge = ""
|
27 |
+
|
28 |
+
def save_game_knowledge(self, new_knowledge):
|
29 |
+
with open(self.game_knowledge_file, 'a', encoding='utf-8') as f:
|
30 |
+
f.write(new_knowledge + "\n")
|
31 |
+
self.game_knowledge += new_knowledge + "\n"
|
32 |
+
|
33 |
+
def generate_response(self, user_input):
|
34 |
+
# Combinar memória, conhecimento do jogo e entrada do usuário
|
35 |
+
context = f"""Conhecimento do Jogo:
|
36 |
+
{self.game_knowledge}
|
37 |
+
|
38 |
+
Histórico de Conversas:
|
39 |
+
{' '.join([f'{m["role"]}: {m["content"]}' for m in self.memory[-5:]])}
|
40 |
+
|
41 |
+
Usuário: {user_input}
|
42 |
+
Assistente:"""
|
43 |
+
|
44 |
+
# Gerar resposta
|
45 |
+
inputs = self.tokenizer(context, return_tensors="pt", max_length=1024, truncation=True)
|
46 |
+
outputs = self.model.generate(
|
47 |
+
inputs["input_ids"],
|
48 |
+
max_length=2048,
|
49 |
+
temperature=0.7,
|
50 |
+
top_p=0.9,
|
51 |
+
pad_token_id=self.tokenizer.eos_token_id
|
52 |
+
)
|
53 |
+
|
54 |
+
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
55 |
+
response = response.split("Assistente:")[-1].strip()
|
56 |
+
|
57 |
+
# Atualizar memória
|
58 |
+
self.memory.append({"role": "user", "content": user_input})
|
59 |
+
self.memory.append({"role": "assistant", "content": response})
|
60 |
+
|
61 |
+
# Manter apenas as últimas N mensagens
|
62 |
+
if len(self.memory) > self.max_memory:
|
63 |
+
self.memory = self.memory[-self.max_memory:]
|
64 |
+
|
65 |
+
return response
|
66 |
+
|
67 |
+
def train_on_new_data(self, training_text):
|
68 |
+
# Salvar dados de treinamento
|
69 |
+
with open("train_data.txt", "w", encoding="utf-8") as f:
|
70 |
+
f.write(training_text)
|
71 |
+
|
72 |
+
# Criar dataset
|
73 |
+
dataset = TextDataset(
|
74 |
+
tokenizer=self.tokenizer,
|
75 |
+
file_path="train_data.txt",
|
76 |
+
block_size=128
|
77 |
+
)
|
78 |
+
|
79 |
+
data_collator = DataCollatorForLanguageModeling(
|
80 |
+
tokenizer=self.tokenizer,
|
81 |
+
mlm=False
|
82 |
+
)
|
83 |
+
|
84 |
+
# Configurar treinamento
|
85 |
+
training_args = TrainingArguments(
|
86 |
+
output_dir="./game_ai_model",
|
87 |
+
overwrite_output_dir=True,
|
88 |
+
num_train_epochs=3,
|
89 |
+
per_device_train_batch_size=4,
|
90 |
+
save_steps=10_000,
|
91 |
+
save_total_limit=2,
|
92 |
+
)
|
93 |
+
|
94 |
+
# Iniciar treinamento
|
95 |
+
trainer = Trainer(
|
96 |
+
model=self.model,
|
97 |
+
args=training_args,
|
98 |
+
data_collator=data_collator,
|
99 |
+
train_dataset=dataset,
|
100 |
+
)
|
101 |
+
|
102 |
+
trainer.train()
|
103 |
+
|
104 |
+
# Salvar como conhecimento do jogo
|
105 |
+
self.save_game_knowledge(training_text)
|
106 |
+
|
107 |
+
return "Treinamento concluído e conhecimento salvo!"
|
108 |
+
|
109 |
+
# Inicializar o assistente
|
110 |
+
assistant = GameAIAssistant()
|
111 |
+
|
112 |
+
# Criar interface Gradio
|
113 |
+
with gr.Blocks() as interface:
|
114 |
+
gr.Markdown("# Assistente de IA para Desenvolvimento de Jogos")
|
115 |
+
|
116 |
+
with gr.Row():
|
117 |
+
with gr.Column():
|
118 |
+
chatbot = gr.Textbox(label="Chat")
|
119 |
+
msg = gr.Textbox(label="Sua mensagem")
|
120 |
+
send = gr.Button("Enviar")
|
121 |
+
|
122 |
+
with gr.Column():
|
123 |
+
training_data = gr.Textbox(label="Dados de Treinamento", lines=10)
|
124 |
+
train = gr.Button("Treinar IA")
|
125 |
+
|
126 |
+
# Funções de callback
|
127 |
+
def chat(message):
|
128 |
+
response = assistant.generate_response(message)
|
129 |
+
return response
|
130 |
+
|
131 |
+
def train_model(text):
|
132 |
+
return assistant.train_on_new_data(text)
|
133 |
+
|
134 |
+
# Conectar componentes
|
135 |
+
send.click(chat, inputs=msg, outputs=chatbot)
|
136 |
+
train.click(train_model, inputs=training_data, outputs=chatbot)
|
137 |
+
|
138 |
+
# Iniciar a interface
|
139 |
+
interface.launch(share=True)
|