VDrontV3-Mini / GUIvdront.py
MishaGGG's picture
Upload 12 files
769a891 verified
Raw
History Blame Contribute Delete
26.7 kB
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()