Starchik1 commited on
Commit
8bf8b74
·
verified ·
1 Parent(s): cd9f358

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +28 -0
main.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForCausalLM, AutoTokenizer
2
+ import torch
3
+
4
+ # Загрузка модели
5
+ model_name = "mistralai/Mistral-7B-Instruct-v0.3"
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
8
+
9
+ # Функция для генерации текста
10
+ def generate_code(prompt, max_length=512):
11
+ inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
12
+ outputs = model.generate(
13
+ inputs["input_ids"],
14
+ max_length=max_length,
15
+ num_return_sequences=1,
16
+ temperature=0.7,
17
+ top_p=0.9,
18
+ do_sample=True,
19
+ )
20
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
21
+
22
+ # Пример запроса
23
+ prompt = """
24
+ Создай простой код игры в стиле 2D RPG с использованием библиотеки Pygame.
25
+ Добавь игрока, монстров, систему здоровья и столкновений.
26
+ """
27
+ game_code = generate_code(prompt)
28
+ print(game_code)