Update app.py

app.py

@@ -1,259 +1,127 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import os
+from collections.abc import Iterator
+from threading import Thread
+
 import gradio as gr
-import time
-from typing import Tuple
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+DESCRIPTION = """\
+# MIXdevAI Llama
+
+MIXdevAI-llama is fine-tuned Russian model based on Llama 3.2 1B Instruct. Model for chating, coding and other! Created by Kolyadual
+"""
+
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
-# Настройки модели
-MODEL_NAME = "Kolyadual/MIXdevAI-llama"
-DEFAULT_MAX_LENGTH = 512
-DEFAULT_TEMPERATURE = 0.7
-DEFAULT_TOP_P = 0.9
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-class ChatBot:
-    def __init__(self):
-        self.model = None
-        self.tokenizer = None
-        self.pipe = None
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.is_loaded = False
-        
-    def load_model(self):
-        """Загрузка модели"""
-        if self.is_loaded:
-            return True
-            
-        try:
-            print("⏳ Загрузка токенизатора...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                MODEL_NAME,
-                trust_remote_code=True
-            )
-            
-            print("⏳ Загрузка модели...")
-            self.model = AutoModelForCausalLM.from_pretrained(
-                MODEL_NAME,
-                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-                device_map="auto" if self.device == "cuda" else None,
-                trust_remote_code=True,
-                low_cpu_mem_usage=True
-            )
-            
-            if self.device == "cpu":
-                self.model = self.model.to(self.device)
-                
-            print("⏳ Создание пайплайна...")
-            self.pipe = pipeline(
-                "text-generation",
-                model=self.model,
-                tokenizer=self.tokenizer,
-                device=0 if self.device == "cuda" else -1
-            )
-            
-            self.is_loaded = True
-            print("✅ Модель успешно загружена!")
-            return True
-            
-        except Exception as e:
-            print(f"❌ Ошибка загрузки модели: {e}")
-            return False
-    
-    def generate_response(self, 
-                         message: str, 
-                         history: list, 
-                         max_length: int,
-                         temperature: float,
-                         top_p: float) -> Tuple[str, list]:
-        """Генерация ответа"""
-        if not self.is_loaded:
-            if not self.load_model():
-                return "Ошибка: модель не загружена", history
-        
-        try:
-            # Форматируем историю для модели
-            prompt = self._format_chat_prompt(message, history)
-            
-            # Генерируем ответ
-            with torch.no_grad():
-                outputs = self.pipe(
-                    prompt,
-                    max_new_tokens=max_length,
-                    temperature=temperature,
-                    top_p=top_p,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id,
-                    repetition_penalty=1.1
-                )
-            
-            # Извлекаем ответ
-            full_response = outputs[0]['generated_text']
-            response = full_response[len(prompt):].strip()
-            
-            # Добавляем в историю
-            history.append((message, response))
-            
-            return "", history
-            
-        except Exception as e:
-            print(f"❌ Ошибка генерации: {e}")
-            return f"Ошибка: {str(e)}", history
-    
-    def _format_chat_prompt(self, message: str, history: list) -> str:
-        """Форматирование промпта из истории чата"""
-        prompt = ""
-        
-        # Добавляем историю
-        for user_msg, assistant_msg in history:
-            prompt += f"### User: {user_msg}\n### Assistant: {assistant_msg}\n"
-        
-        # Добавляем текущее сообщение
-        prompt += f"### User: {message}\n### Assistant: "
-        
-        return prompt
-    
-    def clear_history(self):
-        """Очистка истории"""
-        return []
+model_id = "Kolyadual/MIXdevAI-llama"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+model.eval()
 
-# Создаем экземпляр бота
-chatbot = ChatBot()
 
-def predict(message: str, 
-           history: list, 
-           max_length: int,
-           temperature: float,
-           top_p: float) -> Tuple[str, list]:
-    """Функция для Gradio"""
-    if not message.strip():
-        return "", history
-    
-    # Генерируем ответ
-    _, updated_history = chatbot.generate_response(
-        message, history, max_length, temperature, top_p
+@spaces.GPU(duration=90)
+def generate(
+    message: str,
+    chat_history: list[dict],
+    max_new_tokens: int = 1024,
+    temperature: float = 0.6,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.2,
+) -> Iterator[str]:
+    conversation = [*chat_history, {"role": "user", "content": message}]
+
+    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+    input_ids = input_ids.to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        {"input_ids": input_ids},
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        num_beams=1,
+        repetition_penalty=repetition_penalty,
     )
-    
-    return "", updated_history
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
+
+
+demo = gr.ChatInterface(
+    fn=generate,
+    additional_inputs=[
+        gr.Slider(
+            label="Max new tokens",
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        ),
+        gr.Slider(
+            label="Temperature",
+            minimum=0.1,
+            maximum=4.0,
+            step=0.1,
+            value=0.6,
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=1000,
+            step=1,
+            value=50,
+        ),
+        gr.Slider(
+            label="Repetition penalty",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.2,
+        ),
+    ],
+    stop_btn=None,
+    examples=[
+        ["Привет! Кто ты и кто тебя создал?"],
+        ["Можете вкратце объяснить, что такое язык программирования Python?"],
+        ["Объясните сюжет «Золушки» одним предложением."],
+        ["Сколько часов потребуется человеку, чтобы съесть вертолет?"],
+        ["Напишите статью объемом 100 слов на тему «Преимущества открытого исходного кода в исследованиях в области искусственного интеллекта»."],
+    ],
+    cache_examples=False,
+    type="messages",
+    description=DESCRIPTION,
+    css_paths="style.css",
+    fill_height=True,
+)
 
-def ui():
-    """Создание интерфейса"""
-    with gr.Blocks(
-        title="MIXdevAI-llama Chat",
-        css="""
-        .gradio-container {
-            max-width: 900px;
-            margin: auto;
-        }
-        .chatbot {
-            min-height: 500px;
-        }
-        .footer {
-            text-align: center;
-            margin-top: 20px;
-            color: #666;
-            font-size: 0.9em;
-        }
-        """
-    ) as demo:
-        # Заголовок
-        gr.Markdown("""
-        # 🤖 MIXdevAI-llama Chat Assistant
-        Модель: [Kolyadual/MIXdevAI-llama](https://huggingface.co/Kolyadual/MIXdevAI-llama)
-        """)
-        
-        # Информация о загрузке
-        status = gr.Markdown("Статус: ⏳ Загрузка модели...")
-        
-        # Чат
-        with gr.Row():
-            with gr.Column(scale=3):
-                chatbot_ui = gr.Chatbot(
-                    label="Диалог",
-                    height=500
-                )
-                
-                msg = gr.Textbox(
-                    label="Ваше сообщение",
-                    placeholder="Введите сообщение...",
-                    lines=2,
-                    max_lines=5
-                )
-                
-                with gr.Row():
-                    submit_btn = gr.Button("📤 Отправить", variant="primary")
-                    clear_btn = gr.Button("🗑️ Очистить историю")
-        
-        # Настройки
-        with gr.Accordion("⚙️ Настройки генерации", open=False):
-            with gr.Row():
-                max_length = gr.Slider(
-                    minimum=64,
-                    maximum=2048,
-                    value=DEFAULT_MAX_LENGTH,
-                    step=64,
-                    label="Максимальная длина ответа"
-                )
-                temperature = gr.Slider(
-                    minimum=0.1,
-                    maximum=2.0,
-                    value=DEFAULT_TEMPERATURE,
-                    step=0.1,
-                    label="Температура (креативность)"
-                )
-                top_p = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=DEFAULT_TOP_P,
-                    step=0.05,
-                    label="Top-p (разнообразие)"
-                )
-        
-        # Футер
-        gr.Markdown("""
-        <div class="footer">
-            <p>Модель автоматически загружается при первшем запросе</p>
-            <p>Для работы на CPU потребуется время на загрузку (~5-10 минут)</p>
-        </div>
-        """)
-        
-        # Обработчики событий
-        submit_event = msg.submit(
-            predict,
-            [msg, chatbot_ui, max_length, temperature, top_p],
-            [msg, chatbot_ui]
-        )
-        
-        submit_btn.click(
-            predict,
-            [msg, chatbot_ui, max_length, temperature, top_p],
-            [msg, chatbot_ui]
-        )
-        
-        clear_btn.click(
-            chatbot.clear_history,
-            outputs=chatbot_ui
-        )
-        
-        # Загружаем модель при старте
-        demo.load(
-            chatbot.load_model,
-            outputs=status
-        ).then(
-            lambda: "Статус: ✅ Модель готова к использованию!",
-            outputs=status
-        )
-    
-    return demo
 
 if __name__ == "__main__":
-    # Запускаем приложение
-    demo = ui()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        theme=gr.themes.Soft(
-            primary_hue="blue",
-            secondary_hue="gray"
-        )
-    )
+    demo.launch()