Update app.py

app.py

@@ -1,63 +1,32 @@
-import gradio as gr
 import torch
-import gc
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Чистка памяти перед стартом
-gc.collect()
-
-MODEL_ID = "." 
+model_name = "название_твоей_модели" # Например, "gpt2" или путь к папке
 
 print("🍌 BananaGPT: Попытка загрузки в float16 (Эконом-режим)...")
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# Загружаем токенизатор
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Используем dtype=torch.float16, чтобы веса весили в 2 раза меньше
+# Загружаем модель
+# ВНИМАНИЕ: заменено torch_dtype на dtype, чтобы не было ворнингов
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    device_map="auto", # Теперь, когда accelerate в requirements, это сработает эффективно
-    torch_dtype=torch.float16,
+    model_name,
+    dtype=torch.float16, 
     low_cpu_mem_usage=True
 )
 
-# Красивый интерфейс
-custom_css = """
-footer {visibility: hidden}
-.gradio-container {background-color: #0b1117 !important; color: #e6edf3 !important;}
-.main-title {text-align: center; color: #f1c40f; font-size: 2.5em; font-weight: bold; margin-bottom: 20px;}
-.message.user {border: 1px solid #30363d !important;}
-.message.bot {background-color: #21262d !important; border: 1px solid #30363d !important;}
-"""
-
-def predict(message, history):
-    # Ограничиваем вход, чтобы не вешать процессор
-    inputs = tokenizer(message, return_tensors="pt").to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        inputs,
-        streamer=streamer,
-        max_new_tokens=80,
-        temperature=0.7,
-        do_sample=True,
-    )
-
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    partial_message = ""
-    for new_token in streamer:
-        partial_message += new_token
-        yield partial_message
-
-with gr.Blocks(css=custom_css, title="BananaGPT") as demo:
-    gr.HTML("<div class='main-title'>🍌 BananaGPT</div>")
-    
-    gr.ChatInterface(
-        fn=predict,
-        type="messages",
-    )
-
-if __name__ == "__main__":
-    demo.queue().launch(show_api=False)
+# Переносим на видеокарту, если она есть
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+print(f"✅ Модель успешно загружена на {device}!")
+
+# Тестовый запуск
+prompt = "Привет, BananaGPT!"
+inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+with torch.no_grad():
+    output = model.generate(**inputs, max_new_tokens=50)
+
+print(tokenizer.decode(output[0], skip_special_tokens=True))