asdaswadefswefr commited on
Commit
683a41b
·
verified ·
1 Parent(s): 9e4973c

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +139 -0
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)