Upload 12 files

GUIvdront.py

@@ -0,0 +1,649 @@
+import os
+import json
+import threading
+import pygame
+import torch
+import torch.nn.functional as F
+from transformers import GPT2TokenizerFast, GPT2Config
+from safetensors.torch import load_file
+from model import VDrontModel
+
+# ------------------------------------------------------------
+# Paths / constants
+# ------------------------------------------------------------
+MODEL_DIR = "./VDrontV3-Mini"
+USER_TOKEN = "<|user|>"
+ASSISTANT_TOKEN = "<|assistant|>"
+
+WINDOW_WIDTH = 900
+WINDOW_HEIGHT = 700
+FPS = 60
+
+# ------------------------------------------------------------
+# Theme colors
+# ------------------------------------------------------------
+def get_theme_colors(theme: str):
+    if theme == "dark":
+        return {
+            "background": (25, 25, 30),
+            "surface": (38, 38, 46),
+            "surface_alt": (50, 50, 60),
+            "text": (230, 230, 235),
+            "text_secondary": (160, 160, 170),
+            "accent": (100, 140, 255),
+            "accent_hover": (130, 165, 255),
+            "border": (70, 70, 85),
+            "input_bg": (35, 35, 45),
+            "user_bubble": (70, 100, 200),
+            "user_text": (255, 255, 255),
+            "ai_bubble": (52, 52, 62),
+            "ai_text": (230, 230, 235),
+            "button": (50, 50, 60),
+            "button_hover": (70, 70, 85),
+            "danger": (200, 80, 80),
+            "success": (80, 180, 120),
+        }
+    else:  # light
+        return {
+            "background": (240, 240, 245),
+            "surface": (255, 255, 255),
+            "surface_alt": (230, 230, 235),
+            "text": (30, 30, 35),
+            "text_secondary": (100, 100, 110),
+            "accent": (60, 90, 200),
+            "accent_hover": (90, 120, 230),
+            "border": (200, 200, 210),
+            "input_bg": (245, 245, 250),
+            "user_bubble": (100, 140, 240),
+            "user_text": (255, 255, 255),
+            "ai_bubble": (225, 225, 230),
+            "ai_text": (30, 30, 35),
+            "button": (220, 220, 225),
+            "button_hover": (200, 200, 210),
+            "danger": (200, 80, 80),
+            "success": (80, 180, 120),
+        }
+
+
+# ------------------------------------------------------------
+# Button class
+# ------------------------------------------------------------
+class Button:
+    def __init__(self, rect, text, callback):
+        self.rect = pygame.Rect(rect)
+        self.text = text
+        self.callback = callback
+        self.hovered = False
+
+    def handle_event(self, event):
+        if event.type == pygame.MOUSEMOTION:
+            self.hovered = self.rect.collidepoint(event.pos)
+        elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
+            if self.rect.collidepoint(event.pos):
+                self.callback()
+
+    def draw(self, surface, colors, font):
+        bg = colors["button_hover"] if self.hovered else colors["button"]
+        pygame.draw.rect(surface, bg, self.rect, border_radius=6)
+        pygame.draw.rect(surface, colors["border"], self.rect, width=1, border_radius=6)
+        text_surf = font.render(self.text, True, colors["text"])
+        text_rect = text_surf.get_rect(center=self.rect.center)
+        surface.blit(text_surf, text_rect)
+
+
+# ------------------------------------------------------------
+# Main application
+# ------------------------------------------------------------
+class VDrontLauncher:
+    def __init__(self):
+        pygame.init()
+        self.screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
+        pygame.display.set_caption("VDrontV3-Launcher")
+        self.clock = pygame.time.Clock()
+        self.running = True
+
+        # State
+        self.theme = "dark"
+        self.colors = get_theme_colors(self.theme)
+        self.settings_open = False
+        self.qualitative = False
+        self.input_text = ""
+        self.input_active = True
+        self.messages = []
+        self.scroll_offset = 0
+        self.max_scroll = 0
+        self.generating = False
+        self.gen_thread = None
+        self.gen_done = False
+        self.gen_result = None
+
+        # Generation params (normal mode by default)
+        self.params = {
+            "temperature": 0.45,
+            "max_new_tokens": 256,
+            "repetition_penalty": 1.1,
+            "top_k": 50,
+            "output_version": 0,
+        }
+
+        # Fonts
+        self.font = self._get_font(18)
+        self.font_small = self._get_font(14)
+        self.font_big = self._get_font(22)
+
+        # Load model
+        self._show_loading("Loading model...")
+        self.tokenizer, self.model, self.device = self._load_model()
+        self._show_loading("Ready")
+
+        # UI elements
+        self.top_buttons = []
+        self.send_button = None
+        self._create_buttons()
+
+    # --------------------------------------------------------
+    # Fonts
+    # --------------------------------------------------------
+    @staticmethod
+    def _get_font(size, bold=False):
+        candidates = ["Arial", "DejaVu Sans", "Segoe UI", "Verdana", "Helvetica"]
+        for name in candidates:
+            path = pygame.font.match_font(name, bold=bold)
+            if path:
+                return pygame.font.Font(path, size)
+        return pygame.font.Font(None, size)
+
+    def _show_loading(self, text):
+        self.screen.fill(self.colors["background"])
+        surf = self.font_big.render(text, True, self.colors["text"])
+        rect = surf.get_rect(center=self.screen.get_rect().center)
+        self.screen.blit(surf, rect)
+        pygame.display.flip()
+
+    # --------------------------------------------------------
+    # Model loading
+    # --------------------------------------------------------
+    def _load_model(self):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        tokenizer = GPT2TokenizerFast.from_pretrained(MODEL_DIR)
+        vocab_size = len(tokenizer)
+
+        special_tokens = [USER_TOKEN, ASSISTANT_TOKEN]
+        tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
+
+        with open(os.path.join(MODEL_DIR, "architecture.json")) as f:
+            arch = json.load(f)
+
+        config = GPT2Config(
+            vocab_size=vocab_size,
+            n_embd=arch["n_embd"],
+            n_head=arch["n_head"],
+            n_layer=arch["n_layer"],
+            n_positions=arch["n_positions"],
+            layer_norm_epsilon=1e-5,
+        )
+
+        model = VDrontModel(
+            config=config,
+            expert_start=arch["expert_start"],
+            expert_end=arch["expert_end"],
+            output_index=arch["output_index"],
+            num_experts=arch["num_experts"],
+            num_output_versions=arch["num_output_versions"],
+        )
+
+        state = load_file(os.path.join(MODEL_DIR, "model.safetensors"))
+        model.load_state_dict(state)
+        model.to(device)
+        model.eval()
+
+        # Resize embeddings if tokenizer was extended
+        if model.embed_tokens.num_embeddings < len(tokenizer):
+            old_embed = model.embed_tokens
+            new_embed = torch.nn.Embedding(len(tokenizer), old_embed.embedding_dim).to(device)
+            new_embed.weight.data[:old_embed.num_embeddings] = old_embed.weight.data.to(device)
+            model.embed_tokens = new_embed
+
+            old_lm_head = model.lm_head
+            new_lm_head = torch.nn.Linear(old_lm_head.in_features, len(tokenizer), bias=False).to(device)
+            new_lm_head.weight.data[:old_lm_head.out_features] = old_lm_head.weight.data.to(device)
+            model.lm_head = new_lm_head
+
+            model.config.vocab_size = len(tokenizer)
+
+        return tokenizer, model, device
+
+    # --------------------------------------------------------
+    # UI creation
+    # --------------------------------------------------------
+    def _create_buttons(self):
+        self.theme_button = Button((20, 10, 120, 30), "", self._toggle_theme)
+        self.qualitative_button = Button((150, 10, 140, 30), "", self._toggle_qualitative)
+        self.settings_button = Button((300, 10, 100, 30), "Settings", self._open_settings)
+        self.clear_button = Button((410, 10, 80, 30), "Clear", self._clear_chat)
+        self.top_buttons = [
+            self.theme_button,
+            self.qualitative_button,
+            self.settings_button,
+            self.clear_button,
+        ]
+        self.send_button = Button((WINDOW_WIDTH - 120, WINDOW_HEIGHT - 60, 100, 40), "Send", self._send_message)
+
+    # --------------------------------------------------------
+    # Button callbacks
+    # --------------------------------------------------------
+    def _toggle_theme(self):
+        self.theme = "light" if self.theme == "dark" else "dark"
+        self.colors = get_theme_colors(self.theme)
+
+    def _toggle_qualitative(self):
+        self.qualitative = not self.qualitative
+        if self.qualitative:
+            self.params = {
+                "temperature": 0.3,
+                "max_new_tokens": 512,
+                "repetition_penalty": 1.4,
+                "top_k": 50,
+                "output_version": 1,
+            }
+        else:
+            self.params = {
+                "temperature": 0.45,
+                "max_new_tokens": 256,
+                "repetition_penalty": 1.1,
+                "top_k": 50,
+                "output_version": 0,
+            }
+
+    def _open_settings(self):
+        self.settings_open = True
+
+    def _clear_chat(self):
+        self.messages.clear()
+        self.scroll_offset = 0
+
+    # --------------------------------------------------------
+    # Generation (run in separate thread)
+    # --------------------------------------------------------
+    def _generate_thread(self, prompt):
+        try:
+            self.model.set_output_version(self.params["output_version"])
+            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
+            generated_tokens = []
+            eos_id = self.tokenizer.eos_token_id
+
+            with torch.no_grad():
+                for _ in range(self.params["max_new_tokens"]):
+                    pos = torch.arange(0, input_ids.size(1), device=self.device).unsqueeze(0)
+                    x = self.model.embed_tokens(input_ids) + self.model.embed_positions(pos)
+                    router_logits = self.model.router(x.mean(dim=1))
+                    expert_idx = router_logits.argmax(dim=-1).item()
+                    self.model.set_expert_version(expert_idx)
+
+                    idx_cond = input_ids[:, -self.model.config.n_positions:]
+                    logits, _ = self.model(idx_cond)
+                    logits = logits[:, -1, :] / self.params["temperature"]
+
+                    for token_id in set(input_ids[0].tolist()):
+                        logits[0, token_id] /= self.params["repetition_penalty"]
+
+                    if self.params["top_k"] is not None and self.params["top_k"] > 0:
+                        v, _ = torch.topk(logits, min(self.params["top_k"], logits.size(-1)))
+                        logits[logits < v[:, [-1]]] = -float("Inf")
+
+                    probs = F.softmax(logits, dim=-1)
+                    idx_next = torch.multinomial(probs, num_samples=1)
+                    next_token = idx_next.item()
+
+                    if next_token == eos_id:
+                        break
+
+                    generated_tokens.append(next_token)
+                    input_ids = torch.cat((input_ids, idx_next), dim=1)
+
+            full_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+            self.gen_result = full_text
+        except Exception as e:
+            self.gen_result = f"[Error] {e}"
+        finally:
+            self.gen_done = True
+
+    def _send_message(self):
+        text = self.input_text.strip()
+        if not text or self.generating:
+            return
+
+        self.messages.append({"role": "user", "text": text})
+        self.input_text = ""
+        self.scroll_offset = 0
+
+        prompt = f"{USER_TOKEN}{text}{ASSISTANT_TOKEN}"
+        self.generating = True
+        self.gen_done = False
+        self.gen_result = None
+        self.gen_thread = threading.Thread(target=self._generate_thread, args=(prompt,), daemon=True)
+        self.gen_thread.start()
+
+    # --------------------------------------------------------
+    # Event handling
+    # --------------------------------------------------------
+    def _handle_events(self):
+        for event in pygame.event.get():
+            if event.type == pygame.QUIT:
+                self.running = False
+
+            if self.settings_open:
+                self._handle_settings_event(event)
+            else:
+                self._handle_main_event(event)
+
+    def _handle_main_event(self, event):
+        # Buttons
+        for btn in self.top_buttons:
+            btn.handle_event(event)
+        self.send_button.handle_event(event)
+
+        # Mouse wheel scroll
+        if event.type == pygame.MOUSEWHEEL:
+            self.scroll_offset = max(0, min(self.max_scroll, self.scroll_offset - event.y * 30))
+
+        # Mouse click for input activation
+        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
+            input_rect = pygame.Rect(20, WINDOW_HEIGHT - 60, WINDOW_WIDTH - 140, 40)
+            self.input_active = input_rect.collidepoint(event.pos)
+
+        # Keyboard input
+        if event.type == pygame.KEYDOWN:
+            if event.key == pygame.K_RETURN:
+                self._send_message()
+            elif event.key == pygame.K_BACKSPACE:
+                self.input_text = self.input_text[:-1]
+            elif event.unicode and event.unicode.isprintable():
+                self.input_text += event.unicode
+
+    def _handle_settings_event(self, event):
+        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
+            # Close button / click outside
+            panel_rect = pygame.Rect(
+                (WINDOW_WIDTH - 500) // 2,
+                (WINDOW_HEIGHT - 380) // 2,
+                500,
+                380,
+            )
+            close_rect = pygame.Rect(panel_rect.right - 35, panel_rect.y + 10, 25, 25)
+            if close_rect.collidepoint(event.pos) or not panel_rect.collidepoint(event.pos):
+                self.settings_open = False
+                return
+
+            # Check row controls
+            for key, action, rect in self.settings_controls:
+                if rect.collidepoint(event.pos):
+                    self._adjust_param(key, action)
+                    break
+
+    # --------------------------------------------------------
+    # Settings
+    # --------------------------------------------------------
+    def _adjust_param(self, key, action):
+        row = next((r for r in self.settings_rows if r["key"] == key), None)
+        if not row:
+            return
+
+        if key == "output_version":
+            self.params[key] = 0 if self.params[key] == 1 else 1
+        else:
+            step = row["step"]
+            value = self.params[key]
+            new_value = value + step if action == "plus" else value - step
+            new_value = max(row["min"], min(row["max"], new_value))
+
+            if isinstance(step, int):
+                new_value = int(round(new_value))
+            else:
+                new_value = round(new_value, 2)
+
+            self.params[key] = new_value
+
+        # Manual adjustment means qualitative preset is no longer active
+        self.qualitative = False
+
+    # --------------------------------------------------------
+    # Update
+    # --------------------------------------------------------
+    def _update(self):
+        # Update button labels
+        self.theme_button.text = f"Theme: {'Dark' if self.theme == 'dark' else 'Light'}"
+        self.qualitative_button.text = f"Qualitative: {'ON' if self.qualitative else 'OFF'}"
+
+        # Check generation completion
+        if self.generating and self.gen_done:
+            result = self.gen_result if self.gen_result is not None else "[No response]"
+            self.messages.append({"role": "ai", "text": result})
+            self.generating = False
+            self.gen_done = False
+            self.gen_result = None
+            self.gen_thread = None
+            self.scroll_offset = 0
+
+    # --------------------------------------------------------
+    # Drawing
+    # --------------------------------------------------------
+    def _draw(self):
+        self.screen.fill(self.colors["background"])
+        self._draw_top_bar()
+        self._draw_chat()
+        self._draw_input()
+        if self.generating:
+            self._draw_typing_indicator()
+        if self.settings_open:
+            self._draw_settings()
+        pygame.display.flip()
+
+    def _draw_top_bar(self):
+        for btn in self.top_buttons:
+            btn.draw(self.screen, self.colors, self.font_small)
+
+    def _draw_input(self):
+        input_rect = pygame.Rect(20, WINDOW_HEIGHT - 60, WINDOW_WIDTH - 140, 40)
+        pygame.draw.rect(self.screen, self.colors["input_bg"], input_rect, border_radius=6)
+        pygame.draw.rect(self.screen, self.colors["border"], input_rect, width=1, border_radius=6)
+
+        # Render input text (clipped)
+        text_surf = self.font.render(self.input_text, True, self.colors["text"])
+        clip_rect = input_rect.inflate(-10, -10)
+        self.screen.set_clip(clip_rect)
+        self.screen.blit(text_surf, (input_rect.x + 10, input_rect.y + 8))
+        self.screen.set_clip(None)
+
+        # Blinking cursor
+        if self.input_active and pygame.time.get_ticks() % 1000 < 500:
+            cursor_x = input_rect.x + 10 + text_surf.get_width() + 2
+            if cursor_x < input_rect.right - 10:
+                pygame.draw.line(
+                    self.screen,
+                    self.colors["text"],
+                    (cursor_x, input_rect.y + 8),
+                    (cursor_x, input_rect.y + 32),
+                    2,
+                )
+
+        self.send_button.draw(self.screen, self.colors, self.font)
+
+    def _draw_typing_indicator(self):
+        text = "AI is typing..."
+        surf = self.font_small.render(text, True, self.colors["text_secondary"])
+        rect = surf.get_rect(topleft=(20, WINDOW_HEIGHT - 75))
+        self.screen.blit(surf, rect)
+
+    def _draw_chat(self):
+        chat_rect = pygame.Rect(20, 50, WINDOW_WIDTH - 40, WINDOW_HEIGHT - 130)
+        pygame.draw.rect(self.screen, self.colors["surface"], chat_rect, border_radius=8)
+
+        # Calculate total content height for scrollbar
+        total_height = 0
+        wrapped_cache = []
+        for msg in self.messages:
+            bubble_width = chat_rect.width - 40
+            wrapped = self._wrap_text(msg["text"], self.font, bubble_width - 20)
+            line_height = self.font.get_linesize()
+            bubble_height = line_height * len(wrapped) + 20
+            total_height += bubble_height + 10  # spacing
+            wrapped_cache.append((msg, wrapped, bubble_height))
+        self.max_scroll = max(0, total_height - chat_rect.height)
+        self.scroll_offset = max(0, min(self.scroll_offset, self.max_scroll))
+
+        self.screen.set_clip(chat_rect)
+        y = chat_rect.bottom - 10 + self.scroll_offset
+
+        for msg, wrapped, bubble_height in reversed(wrapped_cache):
+            bubble_rect = pygame.Rect(chat_rect.x + 10, y - bubble_height, chat_rect.width - 40, bubble_height)
+
+            if bubble_rect.bottom < chat_rect.top:
+                break
+
+            if bubble_rect.top <= chat_rect.bottom:
+                if msg["role"] == "user":
+                    bubble_rect.right = chat_rect.right - 10
+                    bg = self.colors["user_bubble"]
+                    fg = self.colors["user_text"]
+                else:
+                    bubble_rect.left = chat_rect.x + 10
+                    bg = self.colors["ai_bubble"]
+                    fg = self.colors["ai_text"]
+
+                pygame.draw.rect(self.screen, bg, bubble_rect, border_radius=10)
+
+                line_height = self.font.get_linesize()
+                text_y = bubble_rect.y + 10
+                for line in wrapped:
+                    line_surf = self.font.render(line, True, fg)
+                    if msg["role"] == "user":
+                        self.screen.blit(line_surf, (bubble_rect.right - 15 - line_surf.get_width(), text_y))
+                    else:
+                        self.screen.blit(line_surf, (bubble_rect.x + 15, text_y))
+                    text_y += line_height
+
+            y = bubble_rect.y - 10
+
+        self.screen.set_clip(None)
+
+        # Scrollbar
+        if total_height > chat_rect.height:
+            scrollbar_height = max(30, int(chat_rect.height * (chat_rect.height / total_height)))
+            scrollbar_y = chat_rect.y + int((chat_rect.height - scrollbar_height) * (self.scroll_offset / self.max_scroll)) if self.max_scroll > 0 else chat_rect.y
+            scrollbar_rect = pygame.Rect(chat_rect.right - 6, scrollbar_y, 4, scrollbar_height)
+            pygame.draw.rect(self.screen, self.colors["border"], scrollbar_rect, border_radius=2)
+
+    def _draw_settings(self):
+        panel_width = 500
+        panel_height = 380
+        panel_x = (WINDOW_WIDTH - panel_width) // 2
+        panel_y = (WINDOW_HEIGHT - panel_height) // 2
+        panel_rect = pygame.Rect(panel_x, panel_y, panel_width, panel_height)
+
+        # Overlay
+        overlay = pygame.Surface((WINDOW_WIDTH, WINDOW_HEIGHT), pygame.SRCALPHA)
+        overlay.fill((0, 0, 0, 128))
+        self.screen.blit(overlay, (0, 0))
+
+        pygame.draw.rect(self.screen, self.colors["surface"], panel_rect, border_radius=12)
+        pygame.draw.rect(self.screen, self.colors["border"], panel_rect, width=2, border_radius=12)
+
+        # Title
+        title_surf = self.font_big.render("Settings", True, self.colors["text"])
+        self.screen.blit(title_surf, (panel_x + 20, panel_y + 15))
+
+        # Close button
+        close_rect = pygame.Rect(panel_rect.right - 35, panel_y + 10, 25, 25)
+        pygame.draw.rect(self.screen, self.colors["button"], close_rect, border_radius=6)
+        pygame.draw.rect(self.screen, self.colors["border"], close_rect, width=1, border_radius=6)
+        close_text = self.font_small.render("X", True, self.colors["text"])
+        self.screen.blit(close_text, close_text.get_rect(center=close_rect.center))
+
+        # Settings rows
+        self.settings_rows = [
+            {"key": "temperature", "label": "Temperature", "min": 0.1, "max": 2.0, "step": 0.05},
+            {"key": "max_new_tokens", "label": "Max Tokens", "min": 32, "max": 1024, "step": 32},
+            {"key": "repetition_penalty", "label": "Repetition Penalty", "min": 0.8, "max": 2.0, "step": 0.1},
+            {"key": "top_k", "label": "Top K", "min": 0, "max": 100, "step": 5},
+            {"key": "output_version", "label": "Output Version", "min": 0, "max": 1, "step": 1},
+        ]
+        self.settings_controls = []
+
+        for i, row in enumerate(self.settings_rows):
+            y = panel_y + 70 + i * 55
+
+            # Label
+            label_surf = self.font.render(row["label"], True, self.colors["text"])
+            self.screen.blit(label_surf, (panel_x + 25, y))
+
+            # Minus button
+            minus_rect = pygame.Rect(panel_x + 310, y, 30, 30)
+            pygame.draw.rect(self.screen, self.colors["button"], minus_rect, border_radius=6)
+            pygame.draw.rect(self.screen, self.colors["border"], minus_rect, width=1, border_radius=6)
+            minus_text = self.font.render("-", True, self.colors["text"])
+            self.screen.blit(minus_text, minus_text.get_rect(center=minus_rect.center))
+            self.settings_controls.append((row["key"], "minus", minus_rect))
+
+            # Value
+            value_surf = self.font.render(str(self.params[row["key"]]), True, self.colors["text"])
+            value_rect = value_surf.get_rect(center=(panel_x + 370, y + 15))
+            self.screen.blit(value_surf, value_rect)
+
+            # Plus button
+            plus_rect = pygame.Rect(panel_x + 410, y, 30, 30)
+            pygame.draw.rect(self.screen, self.colors["button"], plus_rect, border_radius=6)
+            pygame.draw.rect(self.screen, self.colors["border"], plus_rect, width=1, border_radius=6)
+            plus_text = self.font.render("+", True, self.colors["text"])
+            self.screen.blit(plus_text, plus_text.get_rect(center=plus_rect.center))
+            self.settings_controls.append((row["key"], "plus", plus_rect))
+
+    # --------------------------------------------------------
+    # Text wrapping
+    # --------------------------------------------------------
+    def _wrap_text(self, text, font, max_width):
+        words = text.split(" ")
+        lines = []
+        current = ""
+
+        for word in words:
+            test = word if not current else current + " " + word
+            if font.size(test)[0] <= max_width:
+                current = test
+            else:
+                if current:
+                    lines.append(current)
+                    current = word
+                else:
+                    # Very long word, split by characters
+                    while font.size(word)[0] > max_width:
+                        split_idx = len(word)
+                        for i in range(1, len(word)):
+                            if font.size(word[:i])[0] > max_width:
+                                split_idx = i - 1
+                                break
+                        if split_idx == len(word):
+                            break
+                        lines.append(word[:split_idx])
+                        word = word[split_idx:]
+                    current = word
+        if current:
+            lines.append(current)
+        return lines
+
+    # --------------------------------------------------------
+    # Main loop
+    # --------------------------------------------------------
+    def run(self):
+        while self.running:
+            self.clock.tick(FPS)
+            self._handle_events()
+            self._update()
+            self._draw()
+
+        pygame.quit()
+
+
+if __name__ == "__main__":
+    app = VDrontLauncher()
+    app.run()

added_tokens.json

@@ -0,0 +1,4 @@
+{
+  "<|assistant|>": 50258,
+  "<|user|>": 50257
+}

architecture.json

@@ -0,0 +1 @@
+{"expert_start": 0, "expert_end": 3, "output_index": 8, "num_experts": 3, "num_output_versions": 2, "vocab_size": 50259, "n_embd": 768, "n_head": 12, "n_layer": 9, "n_positions": 1024}

config.json

@@ -0,0 +1 @@
+{"vocab_size": 50259, "n_positions": 1024, "n_embd": 768, "n_layer": 9, "n_head": 12, "n_inner": null, "activation_function": "gelu_new", "resid_pdrop": 0.1, "embd_pdrop": 0.1, "attn_pdrop": 0.1, "layer_norm_epsilon": 1e-05, "initializer_range": 0.02, "summary_type": "cls_index", "summary_use_proj": true, "summary_activation": null, "summary_first_dropout": 0.1, "summary_proj_to_labels": true, "scale_attn_weights": true, "use_cache": false, "scale_attn_by_inverse_layer_idx": false, "reorder_and_upcast_attn": false, "bos_token_id": 50256, "eos_token_id": 50256, "return_dict": true, "output_hidden_states": false, "output_attentions": false, "torchscript": false, "torch_dtype": "float32", "use_bfloat16": false, "tf_legacy_loss": false, "pruned_heads": {}, "tie_word_embeddings": true, "chunk_size_feed_forward": 0, "is_encoder_decoder": false, "is_decoder": false, "cross_attention_hidden_size": null, "add_cross_attention": false, "tie_encoder_decoder": false, "max_length": 20, "min_length": 0, "do_sample": false, "early_stopping": false, "num_beams": 1, "num_beam_groups": 1, "diversity_penalty": 0.0, "temperature": 1.0, "top_k": 50, "top_p": 1.0, "typical_p": 1.0, "repetition_penalty": 1.0, "length_penalty": 1.0, "no_repeat_ngram_size": 0, "encoder_no_repeat_ngram_size": 0, "bad_words_ids": null, "num_return_sequences": 1, "output_scores": false, "return_dict_in_generate": false, "forced_bos_token_id": null, "forced_eos_token_id": null, "remove_invalid_values": false, "exponential_decay_length_penalty": null, "suppress_tokens": null, "begin_suppress_tokens": null, "architectures": ["GPT2LMHeadModel"], "finetuning_task": null, "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, "label2id": {"LABEL_0": 0, "LABEL_1": 1}, "tokenizer_class": null, "prefix": null, "pad_token_id": null, "sep_token_id": null, "decoder_start_token_id": null, "task_specific_params": null, "problem_type": null, "_name_or_path": "Base", "_attn_implementation_autoset": true, "transformers_version": "4.46.3", "model_type": "gpt2", "n_ctx": 1024}

model.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPT2Config
+from transformers.models.gpt2.modeling_gpt2 import GPT2Block
+from typing import List
+import copy
+
+class ExpertBlock(nn.Module):
+    def __init__(self, layers: List[nn.Module], num_versions: int, config=None):
+        super().__init__()
+        self.num_versions = num_versions
+        self.config = config
+        self.versions = nn.ModuleList([
+            nn.ModuleList([copy.deepcopy(layer) for layer in layers])
+            for _ in range(num_versions)
+        ])
+        self.active_version = 0
+
+    def set_version(self, idx):
+        self.active_version = idx
+
+    def forward(self, x, **kwargs):
+        for layer in self.versions[self.active_version]:
+            out = layer(x, **kwargs)
+            if isinstance(out, tuple):
+                x = out[0]
+            else:
+                x = out
+        return x
+
+class VDrontModel(nn.Module):
+    def __init__(self, config, expert_start, expert_end, output_index, num_experts, num_output_versions):
+        super().__init__()
+        self.config = config
+        self.num_experts = num_experts
+        self.num_output_versions = num_output_versions
+        self.expert_start = expert_start
+        self.expert_end = expert_end
+        self.output_index = output_index
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.n_embd)
+        self.embed_positions = nn.Embedding(config.n_positions, config.n_embd)
+
+        all_layers = [GPT2Block(config, layer_idx=i) for i in range(config.n_layer)]
+        expert_layers = all_layers[expert_start:expert_end+1]
+        base_layers = all_layers[expert_end+1:output_index]
+        output_layer = all_layers[output_index]
+
+        self.expert_block = ExpertBlock(expert_layers, num_experts, config=config)
+        self.base_blocks = nn.ModuleList(base_layers)
+        self.output_block = ExpertBlock([output_layer], num_output_versions, config=config)
+
+        self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.router = nn.Linear(config.n_embd, num_experts, bias=False)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+
+    def set_expert_version(self, idx):
+        self.expert_block.set_version(idx)
+
+    def set_output_version(self, idx):
+        self.output_block.set_version(idx)
+
+    def forward(self, input_ids, labels=None, return_router_logits=False):
+        pos = torch.arange(0, input_ids.size(1), device=input_ids.device).unsqueeze(0)
+        x = self.embed_tokens(input_ids) + self.embed_positions(pos)
+
+        router_logits = self.router(x.mean(dim=1)) if return_router_logits else None
+
+        x = self.expert_block(x)
+        for block in self.base_blocks:
+            out = block(x)
+            x = out[0] if isinstance(out, tuple) else out
+        x = self.output_block(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            # Защита: всё, что вне [0, vocab_size), заменяем на -100 (игнорируем)
+            labels = torch.where(
+                (labels >= 0) & (labels < self.config.vocab_size),
+                labels,
+                -100
+            )
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                labels.reshape(-1),
+                ignore_index=-100   # ВАЖНО: именно -100, а не -1
+            )
+
+        if return_router_logits:
+            return logits, loss, router_logits
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens, temperature=1.0, top_k=None, dynamic_expert=True):
+        self.eval()
+        for _ in range(max_new_tokens):
+            if dynamic_expert:
+                pos = torch.arange(0, input_ids.size(1), device=input_ids.device).unsqueeze(0)
+                x = self.embed_tokens(input_ids) + self.embed_positions(pos)
+                router_logits = self.router(x.mean(dim=1))
+                expert_idx = router_logits.argmax(dim=-1).item()
+                self.set_expert_version(expert_idx)
+
+            idx_cond = input_ids[:, -self.config.n_positions:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat((input_ids, idx_next), dim=1)
+        return input_ids

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37a135d8409f26cd4804bbb49a65dd0feba32814751ef36c2f935ad0b2ff6835
+size 822303144

special_tokens_map.json

@@ -0,0 +1,53 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<|user|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|assistant|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,41 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50257": {
+      "content": "<|user|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50258": {
+      "content": "<|assistant|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|user|>",
+    "<|assistant|>",
+    "<|endoftext|>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

use.py

@@ -0,0 +1,123 @@
+# use.py
+import os
+import json
+import torch
+import torch.nn.functional as F
+from transformers import GPT2TokenizerFast, GPT2Config
+from safetensors.torch import load_file
+from model import VDrontModel
+
+CONFIG = {
+    "model_dir": "./VDrontV3-Mini",
+    "temperature": 0.4,
+    "top_k": 50,
+    "max_new_tokens": 200,
+    "repetition_penalty": 1.2,
+    "user_token": "<|user|>",
+    "assistant_token": "<|assistant|>",
+}
+
+def format_prompt(user_input):
+    return f"{CONFIG['user_token']}{user_input}{CONFIG['assistant_token']}"
+
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    tokenizer = GPT2TokenizerFast.from_pretrained(CONFIG['model_dir'])
+    vocab_size = len(tokenizer)
+
+    special_tokens = [CONFIG['user_token'], CONFIG['assistant_token']]
+    tokenizer.add_special_tokens({'additional_special_tokens': special_tokens})
+
+    with open(os.path.join(CONFIG['model_dir'], 'architecture.json')) as f:
+        arch = json.load(f)
+
+    config = GPT2Config(
+        vocab_size=vocab_size,
+        n_embd=arch['n_embd'],
+        n_head=arch['n_head'],
+        n_layer=arch['n_layer'],
+        n_positions=arch['n_positions'],
+        layer_norm_epsilon=1e-5,
+    )
+
+    model = VDrontModel(
+        config=config,
+        expert_start=arch['expert_start'],
+        expert_end=arch['expert_end'],
+        output_index=arch['output_index'],
+        num_experts=arch['num_experts'],
+        num_output_versions=arch['num_output_versions'],
+    )
+    state = load_file(os.path.join(CONFIG['model_dir'], 'model.safetensors'))
+    model.load_state_dict(state)
+    model.to(device)
+    model.eval()
+
+    if model.embed_tokens.num_embeddings < len(tokenizer):
+        old_embed = model.embed_tokens
+        new_embed = torch.nn.Embedding(len(tokenizer), old_embed.embedding_dim).to(device)
+        new_embed.weight.data[:old_embed.num_embeddings] = old_embed.weight.data.to(device)
+        model.embed_tokens = new_embed
+
+        old_lm_head = model.lm_head
+        new_lm_head = torch.nn.Linear(old_lm_head.in_features, len(tokenizer), bias=False).to(device)
+        new_lm_head.weight.data[:old_lm_head.out_features] = old_lm_head.weight.data.to(device)
+        model.lm_head = new_lm_head
+
+        model.config.vocab_size = len(tokenizer)
+
+    while True:
+        try:
+            output_ver = int(input("Mode (0 - base (bad, little answer), 1 - qualitative (normal, medium answer): "))
+            if output_ver in [0, 1]:
+                model.set_output_version(output_ver)
+                break
+        except ValueError:
+            pass
+
+    print("Chat is ready. Type 'exit' to quit.")
+
+    while True:
+        user_input = input("You: ")
+        if user_input.lower() in ['exit', 'quit']:
+            break
+
+        prompt = format_prompt(user_input)
+        input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
+        generated_tokens = []
+        eos_id = tokenizer.eos_token_id
+
+        with torch.no_grad():
+            for _ in range(CONFIG['max_new_tokens']):
+                pos = torch.arange(0, input_ids.size(1), device=device).unsqueeze(0)
+                x = model.embed_tokens(input_ids) + model.embed_positions(pos)
+                router_logits = model.router(x.mean(dim=1))
+                expert_idx = router_logits.argmax(dim=-1).item()
+                model.set_expert_version(expert_idx)
+
+                idx_cond = input_ids[:, -model.config.n_positions:]
+                logits, _ = model(idx_cond)
+                logits = logits[:, -1, :] / CONFIG['temperature']
+
+                for token_id in set(input_ids[0].tolist()):
+                    logits[0, token_id] /= CONFIG['repetition_penalty']
+
+                if CONFIG['top_k'] is not None:
+                    v, _ = torch.topk(logits, min(CONFIG['top_k'], logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = -float('Inf')
+
+                probs = F.softmax(logits, dim=-1)
+                idx_next = torch.multinomial(probs, num_samples=1)
+                next_token = idx_next.item()
+
+                if next_token == eos_id:
+                    break
+
+                generated_tokens.append(next_token)
+                input_ids = torch.cat((input_ids, idx_next), dim=1)
+
+        full_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        print(f"AI: {full_text}")
+
+if __name__ == '__main__':
+    main()