Update generate_game.py

generate_game.py

@@ -1,28 +1,10 @@
+import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
-# Загрузка модели
+# Загрузка токена из переменной окружения
+hf_token = os.getenv("HUGGING_FACE_TOKEN")
 model_name = "mistralai/Mistral-7B-Instruct-v0.3"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-# Функция для генерации текста
-def generate_code(prompt, max_length=512):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(
-        inputs["input_ids"],
-        max_length=max_length,
-        num_return_sequences=1,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
 
-# Пример запроса
-prompt = """
-Создай простой код игры в стиле 2D RPG с использованием библиотеки Pygame. 
-Добавь игрока, монстров, систему здоровья и столкновений.
-"""
-game_code = generate_code(prompt)
-print(game_code)