Final update including testing, final model, new data sources, updated notebook for mlp and logistic regression and final keyboard code

app/constants.py

@@ -0,0 +1,6 @@
+CONTEXT_SIZE = 10
+ALPHABET = list("abcdefghijklmnopqrstuvwxyz")
+VOCAB_SIZE = len(ALPHABET) + 1  # +1 for unknown
+UNKNOWN_IDX = VOCAB_SIZE - 1
+
+SINGLE_KEY_PROBABILITY = 1 / 26

app/core/config.py

@@ -14,7 +14,7 @@ class Settings(BaseSettings):
 
     APP_NAME: str = "Omega"
     VERBOSITY: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO"
-    MODEL_PATH: FilePath = Path("./model/predictor.pth")
+    MODEL_PATH: FilePath = Path("model.onnx")
 
 
 settings = Settings()

app/keyboard.py

@@ -7,7 +7,7 @@ from app.ui.view import KeyboardView
 
 from app.utils import setup_logger
 
-class Keyboard():
+class KeyboardApplication():
     def __init__(self):
         setup_logger()
         self.logger = logging.getLogger(__name__)

app/predictor.py

@@ -0,0 +1,56 @@
+import onnxruntime
+import numpy as np
+import sys
+
+from app.core.config import settings
+from app.constants import ALPHABET, CONTEXT_SIZE, SINGLE_KEY_PROBABILITY
+from app.utils import tokenize
+
+import logging
+
+class KeyboardPredictor():
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+        self.buffer = ""
+        try:
+            self.session = onnxruntime.InferenceSession(settings.MODEL_PATH)
+        except onnxruntime.capi.onnxruntime_pybind11_state.InvalidProtobuf as e:
+            self.logger.critical("Invalid model found. Exiting...")
+            sys.exit(1)
+
+
+    def reset(self):
+        self.logger.info("Clearing prediction buffer")
+        self.buffer = ""
+
+    def update(self, char: str):
+        self.logger.info("Adding new char '%s'", char)
+        # Remove the oldest character
+        # and append the current one
+        if len(self.buffer) == CONTEXT_SIZE:
+            self.buffer = self.buffer[1:]
+        self.buffer += char.lower()
+
+    def __get_prediction(self) -> dict[str, float]:
+        """Returns all probabilities distribution over the next character."""
+        tokenized = [tokenize("".join(self.buffer))]  # shape: (1, context_size)
+        input_array = np.array(tokenized, dtype=np.int64)
+
+        outputs = self.session.run(None, {"input": input_array})
+        logits = outputs[0][0]  # shape: (VOCAB_SIZE,)
+
+        exp_logits = np.exp(logits - np.max(logits))  # numerical stability
+        probs = exp_logits / exp_logits.sum()
+
+        return {char: probs[i] for i, char in enumerate(ALPHABET + ["?"])}
+
+    def get_predictions(self) -> dict[str, float]:
+        """
+            Returns keys that are most likely to be pressed next.
+        """
+        self.logger.debug("Getting fresh predictions")
+        DROPOUT_PROBABILITY = 2 * SINGLE_KEY_PROBABILITY
+
+        probs = self.__get_prediction()
+        probs = {char: prob for char, prob in probs.items() if prob > DROPOUT_PROBABILITY}
+        return probs

app/ui/current_word.py

@@ -0,0 +1,54 @@
+from PySide6.QtWidgets import QGraphicsRectItem, QGraphicsSimpleTextItem
+from PySide6.QtGui import QFont, QFontMetricsF, QColor, QBrush
+
+
+BACKGROUND_COLOR = QColor("white")
+TEXT_SCALE_FACTOR = 0.4  # 40% of the width or height
+Z_INDEX = 10
+
+
+class CurrentWordBox(QGraphicsRectItem):
+    def __init__(self):
+        super().__init__()
+
+        self.word = ""
+        self.text = QGraphicsSimpleTextItem("", self)
+        self.text.setFlag(QGraphicsSimpleTextItem.ItemIgnoresTransformations, True)
+
+        self.setBrush(QBrush(BACKGROUND_COLOR))
+        self.setZValue(Z_INDEX)
+
+        self.width = 0
+        self.height = 0
+
+    def set_geometry(self, x: float, y: float, width: float, height: float):
+        self.width = width
+        self.height = height
+        self.setRect(0, 0, width, height)
+        self.setPos(x, y)
+        self.update_label_font()
+
+    def update_word(self, next_char: str):
+        self.word += next_char
+        self.text.setText(self.word)
+        self.update_label_font()
+
+    def clear(self):
+        self.word = ""
+        self.text.setText("")
+        self.update_label_font()
+
+    def update_label_font(self):
+        min_dimension = min(self.width, self.height)
+        font_size = min_dimension * TEXT_SCALE_FACTOR
+
+        font = QFont()
+        font.setPointSizeF(font_size)
+        self.text.setFont(font)
+
+        metrics = QFontMetricsF(font)
+        text_rect = metrics.boundingRect(self.word)
+
+        text_x = (self.width - text_rect.width()) / 2
+        text_y = (self.height - text_rect.height()) / 2
+        self.text.setPos(text_x, text_y)

app/ui/key.py

@@ -1,23 +1,45 @@
+import logging
+
 from PySide6.QtWidgets import (
     QGraphicsRectItem, QGraphicsSimpleTextItem, QGraphicsItem
 )
-from PySide6.QtGui import QBrush, QColor, QFont, QFontMetricsF
+from PySide6.QtGui import QBrush, QColor, QFont, QFontMetricsF, QPen, Qt
+from PySide6.QtCore import Signal, QObject, QTimer
 
 
-class KeyItem(QGraphicsRectItem):
-    def __init__(self, label):
-        super().__init__()
+NORMAL_COLOR = QColor("lightgray")
+CLICKED_COLOR = QColor("darkgray")
+TEXT_SCALE_FACTOR = 0.4
+CLICK_HIGHLIGHT_DURATION_MS = 100
+Z_INDEX_MULTIPLIER = 2.0
+
+class KeyItem(QGraphicsRectItem, QObject):
+    clicked = Signal(str)
+
+    def __init__(self, label: str):
+        QObject.__init__(self)
+        QGraphicsRectItem.__init__(self)
+
+        self.logger = logging.getLogger(__name__)
+        
+
         self.label = label
+        self.width = 0.0
+        self.height = 0.0
+        self.scale_factor = 1.0
+
+        self._reset_timer = QTimer()
+        self._reset_timer.setSingleShot(True)
+        self._reset_timer.timeout.connect(self.__reset_color)
+
         self.text = QGraphicsSimpleTextItem(label, self)
         self.text.setFlag(QGraphicsItem.ItemIgnoresTransformations, True)
 
-        self.setBrush(QBrush(QColor("lightgray")))
+        self.__normal_brush = QBrush(NORMAL_COLOR)
+        self.__click_brush = QBrush(CLICKED_COLOR)
+        self.setBrush(self.__normal_brush)
 
-        self.scale_factor = 1.0
-        self.width = 0
-        self.height = 0
-
-    def set_geometry(self, x, y, width, height):
+    def set_geometry(self, x: float, y: float, width: float, height: float):
         self.width = width
         self.height = height
 
@@ -28,25 +50,34 @@ class KeyItem(QGraphicsRectItem):
         self.update_label_font()
 
     def update_label_font(self):
-        # Dynamically size font
-        font_size = min(self.width, self.height) * 0.4
+        min_dimension = min(self.width, self.height)
+        font_size = min_dimension * TEXT_SCALE_FACTOR
+
         font = QFont()
         font.setPointSizeF(font_size)
         self.text.setFont(font)
 
-        # Use font metrics for accurate bounding
         metrics = QFontMetricsF(font)
         text_rect = metrics.boundingRect(self.label)
 
         text_x = (self.width - text_rect.width()) / 2
-        text_y = (self.height - text_rect.height()) / 2 # + metrics.ascent() - text_rect.height()
-
+        text_y = (self.height - text_rect.height()) / 2
         self.text.setPos(text_x, text_y)
 
-    def set_scale_factor(self, scale):
-        self.scale_factor = scale
-        self.setZValue(scale)
+    def set_scale_factor(self, scale: float):
+        scaled_z = scale * Z_INDEX_MULTIPLIER
+        self.scale_factor = scaled_z
+        self.setZValue(scaled_z)
         self.setScale(scale)
 
     def mousePressEvent(self, q_mouse_event):
-        print(self.label, "was clicked")
+        self.logger.info("%s was clicked", self.label)
+        self.clicked.emit(self.label)
+        self.__onclick_color()
+
+    def __onclick_color(self):
+        self.setBrush(self.__click_brush)
+        self._reset_timer.start(CLICK_HIGHLIGHT_DURATION_MS)
+
+    def __reset_color(self):
+        self.setBrush(self.__normal_brush)

app/ui/keyboard.py

@@ -0,0 +1,70 @@
+from PySide6.QtCore import QSize, QRectF
+from app.ui.key import KeyItem
+
+
+class KeyboardLayout():
+    def __init__(self, scene):
+        self.scene = scene
+        self.keys = {}
+        self.letter_rows = [
+            "QWERTZUIOP",
+            "ASDFGHJKL",
+            "YXCVBNM",
+        ]
+        self.space_key = KeyItem("Space")
+        self.space_key.clicked.connect(lambda: self.on_key_clicked("Space"))
+        self.scene.addItem(self.space_key)
+        self.create_keys()
+
+        # callback to be set by scene
+        self.key_pressed_callback = None
+
+    def create_keys(self):
+        for row in self.letter_rows:
+            for char in row:
+                key = KeyItem(char)
+                key.clicked.connect(lambda c=char: self.on_key_clicked(c))
+                self.scene.addItem(key)
+                self.keys[char] = key
+
+    def layout_keys(self, view_size: QSize, top_offset: float = 0.0):
+        view_width = view_size.width()
+        view_height = view_size.height()
+
+        padding = 10.0
+        spacing = 6.0
+        rows = self.letter_rows + [" "]
+        row_count = len(rows)
+        max_keys_in_row = max(len(row) for row in rows)
+
+        total_spacing_x = (max_keys_in_row - 1) * spacing
+        total_spacing_y = (row_count - 1) * spacing
+
+        available_width = view_width - 2 * padding - total_spacing_x
+        available_height = view_height - top_offset - 2 * padding - total_spacing_y
+
+        key_width = available_width / max_keys_in_row
+        key_height = available_height / row_count
+
+        y = top_offset + padding
+        for row in self.letter_rows:
+            x = padding
+            for char in row:
+                key = self.keys[char]
+                key.set_geometry(x, y, key_width, key_height)
+                x += key_width + spacing
+            y += key_height + spacing
+
+        space_width = key_width * 7 + spacing * 6
+        self.space_key.set_geometry(padding, y, space_width, key_height)
+
+    def set_scale_factors(self, predictions: dict):
+        for key in self.keys.values():
+            key.set_scale_factor(1.0)
+        for key, probability in predictions.items():
+            if key.upper() in self.keys:
+                self.keys[key.upper()].set_scale_factor(1 + probability)
+
+    def on_key_clicked(self, label: str):
+        if self.key_pressed_callback:
+            self.key_pressed_callback(label)

app/ui/scene.py

@@ -1,79 +1,52 @@
-from PySide6.QtWidgets import (
-    QGraphicsScene
-)
-from PySide6.QtGui import QBrush, QColor
-from PySide6.QtCore import QTimer, QRectF, QSize
-import random
+from PySide6.QtWidgets import QGraphicsScene
+from PySide6.QtCore import QRectF, QSize
+
+from app.predictor import KeyboardPredictor
+from app.ui.current_word import CurrentWordBox
+from app.ui.keyboard import KeyboardLayout
+
+
+WORD_BOX_HEIGHT_MODIFIER = 0.15
+WORD_BOX_PADDING_ALL_SIDES = 10
+WORD_BOX_MARGIN_BOTTOM = 20
 
-from app.ui.key import KeyItem
 class KeyboardScene(QGraphicsScene):
     def __init__(self):
         super().__init__()
-        self.keys = {}
-        self.letter_rows = [
-            "QWERTYUIOP",
-            "ASDFGHJKL",
-            "ZXCVBNM",
-        ]
-        self.space_key = KeyItem("Space")
-        self.addItem(self.space_key)
-        self.timer = QTimer()
-        self.timer.timeout.connect(self.simulate_prediction)
-        self.timer.start(2000)
 
-        self.create_keys()
+        self.word_box = CurrentWordBox()
+        self.addItem(self.word_box)
 
-    def create_keys(self):
-        for row in self.letter_rows:
-            for char in row:
-                key = KeyItem(char)
-                self.addItem(key)
-                self.keys[char] = key
+        self.predictor = KeyboardPredictor()
+
+        self.keyboard = KeyboardLayout(self)
+        self.keyboard.key_pressed_callback = self.on_key_clicked
 
     def layout_keys(self, view_size: QSize):
         view_width = view_size.width()
         view_height = view_size.height()
 
-        padding = 10.0
-        spacing = 6.0
-        max_scale = 2.0
+        word_box_height = view_height * WORD_BOX_HEIGHT_MODIFIER
+        word_box_x = WORD_BOX_PADDING_ALL_SIDES
+        word_box_y = WORD_BOX_PADDING_ALL_SIDES
+        word_box_width = view_width - WORD_BOX_MARGIN_BOTTOM
+        word_box_height_adjusted = word_box_height - WORD_BOX_MARGIN_BOTTOM
 
-        rows = self.letter_rows + [" "]
-        row_count = len(rows)
-        max_keys_in_row = max(len(row) for row in rows)
-
-        total_spacing_x = (max_keys_in_row - 1) * spacing
-        total_spacing_y = (row_count - 1) * spacing
-
-        available_width = view_width - 2 * padding - total_spacing_x
-        available_height = view_height - 2 * padding - total_spacing_y
-
-        key_width = available_width / max_keys_in_row
-        key_height = available_height / row_count
-
-        y = padding
-        for row in self.letter_rows:
-            x = padding
-            for char in row:
-                key = self.keys[char]
-                key.set_geometry(x, y, key_width, key_height)
-                x += key_width + spacing
-            y += key_height + spacing
-
-        # Space key layout
-        space_width = key_width * 7 + spacing * 6
-        self.space_key.set_geometry(padding, y, space_width, key_height)
+        self.word_box.set_geometry(
+            word_box_x, word_box_y, word_box_width, word_box_height_adjusted
+        )
 
+        self.keyboard.layout_keys(view_size, top_offset=word_box_height)
         self.setSceneRect(QRectF(0, 0, view_width, view_height))
 
-    def simulate_prediction(self):
-        most_likely = random.choice(list(self.keys.keys()))
-        print(f"[Prediction] Most likely: {most_likely}")
+    def on_key_clicked(self, label: str):
+        if label == "Space":
+            self.word_box.clear()
+            self.predictor.reset()
+            self.keyboard.set_scale_factors({})
+            return
 
-        for char, key in self.keys.items():
-            if char == most_likely:
-                key.set_scale_factor(1.8)
-                key.setBrush(QBrush(QColor("orange")))
-            else:
-                key.set_scale_factor(1.0)
-                key.setBrush(QBrush(QColor("lightgray")))
+        self.word_box.update_word(label)
+        self.predictor.update(label)
+        predictions = self.predictor.get_predictions()
+        self.keyboard.set_scale_factors(predictions)

app/ui/view.py

@@ -5,6 +5,7 @@ from PySide6.QtGui import QPainter
 from PySide6.QtCore import Qt
 
 from .scene import KeyboardScene
+from app.core.config import settings
 
 class KeyboardView(QGraphicsView):
     def __init__(self):
@@ -12,9 +13,9 @@ class KeyboardView(QGraphicsView):
         self.scene = KeyboardScene()
         self.setScene(self.scene)
         self.setRenderHint(QPainter.Antialiasing)
-        self.setWindowTitle("Dynamic Keyboard")
+        self.setWindowTitle(settings.APP_NAME)
         self.setAlignment(Qt.AlignLeft | Qt.AlignTop)
-        self.setMinimumSize(600, 200)  # Sensible default
+        self.setMinimumSize(600, 300)
 
     def resizeEvent(self, event):
         super().resizeEvent(event)

app/utils.py

@@ -3,6 +3,11 @@ import logging
 
 from app.core.config import settings
 
+from app.constants import CONTEXT_SIZE, UNKNOWN_IDX, ALPHABET
+
+char_to_index = {char: idx for idx, char in enumerate(ALPHABET)}
+
+
 def setup_logger():
     logger = logging.getLogger()
     
@@ -50,4 +55,14 @@ if sys.platform == "win32":
             hwnd, win32con.HWND_TOPMOST,
             0, 0, 0, 0,
             win32con.SWP_NOMOVE | win32con.SWP_NOSIZE | win32con.SWP_NOACTIVATE
-        )
+        )
+
+
+def tokenize(text: str) -> list[int]:
+    """Convert last CONTEXT_SIZE chars to integer indices."""
+    text = text.lower()[-CONTEXT_SIZE:]  # trim to context length
+    padded = [' '] * (CONTEXT_SIZE - len(text)) + list(text)
+    return [
+        char_to_index.get(c, UNKNOWN_IDX)
+        for c in padded
+    ]