text_test.c

/*
 * Text Sentiment Classification Test
 * Demonstrates the neural network library with text instead of images.
 *
 * Encoding: Character frequency (a-z → 26-dimensional vector)
 * Task: Classify text as POSITIVE (0) or NEGATIVE (1)
 */

#define _CRT_SECURE_NO_WARNINGS

#include "arena.h"
#include "base.h"
#include "prng.h"


#include "arena.c"
#include "prng.c"

// ============== Matrix Operations (copied from main.c) ==============

typedef struct {
  u32 rows, cols;
  f32 *data;
} matrix;

matrix *mat_create(mem_arena *arena, u32 rows, u32 cols) {
  matrix *mat = PUSH_STRUCT(arena, matrix);
  mat->rows = rows;
  mat->cols = cols;
  mat->data = PUSH_ARRAY(arena, f32, (u64)rows * cols);
  return mat;
}

void mat_clear(matrix *mat) {
  memset(mat->data, 0, sizeof(f32) * (u64)mat->rows * mat->cols);
}

void mat_fill(matrix *mat, f32 x) {
  u64 size = (u64)mat->rows * mat->cols;
  for (u64 i = 0; i < size; i++) {
    mat->data[i] = x;
  }
}

void mat_fill_rand(matrix *mat, f32 lower, f32 upper) {
  u64 size = (u64)mat->rows * mat->cols;
  for (u64 i = 0; i < size; i++) {
    mat->data[i] = prng_randf() * (upper - lower) + lower;
  }
}

void mat_scale(matrix *mat, f32 scale) {
  u64 size = (u64)mat->rows * mat->cols;
  for (u64 i = 0; i < size; i++) {
    mat->data[i] *= scale;
  }
}

f32 mat_sum(matrix *mat) {
  u64 size = (u64)mat->rows * mat->cols;
  f32 sum = 0.0f;
  for (u64 i = 0; i < size; i++) {
    sum += mat->data[i];
  }
  return sum;
}

u64 mat_argmax(matrix *mat) {
  u64 size = (u64)mat->rows * mat->cols;
  u64 max_i = 0;
  for (u64 i = 0; i < size; i++) {
    if (mat->data[i] > mat->data[max_i]) {
      max_i = i;
    }
  }
  return max_i;
}

void mat_add(matrix *out, const matrix *a, const matrix *b) {
  u64 size = (u64)out->rows * out->cols;
  for (u64 i = 0; i < size; i++) {
    out->data[i] = a->data[i] + b->data[i];
  }
}

void mat_sub(matrix *out, const matrix *a, const matrix *b) {
  u64 size = (u64)out->rows * out->cols;
  for (u64 i = 0; i < size; i++) {
    out->data[i] = a->data[i] - b->data[i];
  }
}

void mat_mul(matrix *out, const matrix *a, const matrix *b, b8 zero_out,
             b8 transpose_a, b8 transpose_b) {
  u32 a_rows = transpose_a ? a->cols : a->rows;
  u32 a_cols = transpose_a ? a->rows : a->cols;
  u32 b_cols = transpose_b ? b->rows : b->cols;

  if (zero_out) {
    mat_clear(out);
  }

  for (u64 i = 0; i < a_rows; i++) {
    for (u64 k = 0; k < a_cols; k++) {
      f32 a_val =
          transpose_a ? a->data[i + k * a->cols] : a->data[k + i * a->cols];
      for (u64 j = 0; j < b_cols; j++) {
        f32 b_val =
            transpose_b ? b->data[k + j * b->cols] : b->data[j + k * b->cols];
        out->data[j + i * out->cols] += a_val * b_val;
      }
    }
  }
}

void mat_relu(matrix *out, const matrix *in) {
  u64 size = (u64)out->rows * out->cols;
  for (u64 i = 0; i < size; i++) {
    out->data[i] = in->data[i] > 0 ? in->data[i] : 0;
  }
}

void mat_softmax(matrix *out, const matrix *in) {
  u64 size = (u64)out->rows * out->cols;
  f32 sum = 0.0f;
  for (u64 i = 0; i < size; i++) {
    out->data[i] = expf(in->data[i]);
    sum += out->data[i];
  }
  mat_scale(out, 1.0f / sum);
}

// ============== Text Encoding ==============

#define INPUT_SIZE 26 // a-z character frequencies
#define HIDDEN_SIZE 8
#define OUTPUT_SIZE 2 // positive=0, negative=1

void encode_text(const char *text, f32 *output) {
  memset(output, 0, sizeof(f32) * INPUT_SIZE);

  u32 total = 0;
  for (const char *p = text; *p; p++) {
    char c = *p;
    if (c >= 'A' && c <= 'Z')
      c = c - 'A' + 'a';
    if (c >= 'a' && c <= 'z') {
      output[c - 'a'] += 1.0f;
      total++;
    }
  }

  // Normalize
  if (total > 0) {
    for (u32 i = 0; i < INPUT_SIZE; i++) {
      output[i] /= (f32)total;
    }
  }
}

// ============== Training Data ==============

typedef struct {
  const char *text;
  u32 label; // 0=positive, 1=negative
} sample;

sample training_data[] = {
    // Positive (label=0)
    {"I love this product", 0},
    {"This is amazing", 0},
    {"Great quality", 0},
    {"Excellent work", 0},
    {"Very happy with this", 0},
    {"Wonderful experience", 0},
    {"Best purchase ever", 0},
    {"Highly recommend", 0},
    {"Fantastic job", 0},
    {"Perfect solution", 0},
    {"Love it so much", 0},
    {"Really impressed", 0},
    {"Super good quality", 0},
    {"Awesome product", 0},
    {"Very satisfied", 0},

    // Negative (label=1)
    {"This is terrible", 1},
    {"Very bad quality", 1},
    {"I hate this", 1},
    {"Worst purchase", 1},
    {"Not recommended", 1},
    {"Very disappointed", 1},
    {"Awful experience", 1},
    {"Poor quality", 1},
    {"Complete waste", 1},
    {"Really bad", 1},
    {"Horrible product", 1},
    {"Never buying again", 1},
    {"Total disaster", 1},
    {"Broken on arrival", 1},
    {"Extremely unhappy", 1},
};

#define NUM_SAMPLES (sizeof(training_data) / sizeof(training_data[0]))

// ============== Simple Neural Network ==============

typedef struct {
  matrix *W1, *b1, *W2, *b2;     // Parameters
  matrix *z1, *a1, *z2, *out;    // Forward values
  matrix *dW1, *db1, *dW2, *db2; // Gradients
  matrix *da1, *dz1, *dz2;       // Intermediate gradients
  matrix *input, *target;
} network;

network *create_network(mem_arena *arena) {
  network *net = PUSH_STRUCT(arena, network);

  // Parameters
  net->W1 = mat_create(arena, HIDDEN_SIZE, INPUT_SIZE);
  net->b1 = mat_create(arena, HIDDEN_SIZE, 1);
  net->W2 = mat_create(arena, OUTPUT_SIZE, HIDDEN_SIZE);
  net->b2 = mat_create(arena, OUTPUT_SIZE, 1);

  // Xavier initialization
  f32 bound1 = sqrtf(6.0f / (INPUT_SIZE + HIDDEN_SIZE));
  f32 bound2 = sqrtf(6.0f / (HIDDEN_SIZE + OUTPUT_SIZE));
  mat_fill_rand(net->W1, -bound1, bound1);
  mat_fill_rand(net->W2, -bound2, bound2);

  // Forward pass storage
  net->z1 = mat_create(arena, HIDDEN_SIZE, 1);
  net->a1 = mat_create(arena, HIDDEN_SIZE, 1);
  net->z2 = mat_create(arena, OUTPUT_SIZE, 1);
  net->out = mat_create(arena, OUTPUT_SIZE, 1);

  // Gradients
  net->dW1 = mat_create(arena, HIDDEN_SIZE, INPUT_SIZE);
  net->db1 = mat_create(arena, HIDDEN_SIZE, 1);
  net->dW2 = mat_create(arena, OUTPUT_SIZE, HIDDEN_SIZE);
  net->db2 = mat_create(arena, OUTPUT_SIZE, 1);
  net->da1 = mat_create(arena, HIDDEN_SIZE, 1);
  net->dz1 = mat_create(arena, HIDDEN_SIZE, 1);
  net->dz2 = mat_create(arena, OUTPUT_SIZE, 1);

  // Input/target
  net->input = mat_create(arena, INPUT_SIZE, 1);
  net->target = mat_create(arena, OUTPUT_SIZE, 1);

  return net;
}

void forward(network *net) {
  // Layer 1: z1 = W1 * input + b1, a1 = relu(z1)
  mat_mul(net->z1, net->W1, net->input, 1, 0, 0);
  mat_add(net->z1, net->z1, net->b1);
  mat_relu(net->a1, net->z1);

  // Layer 2: z2 = W2 * a1 + b2, out = softmax(z2)
  mat_mul(net->z2, net->W2, net->a1, 1, 0, 0);
  mat_add(net->z2, net->z2, net->b2);
  mat_softmax(net->out, net->z2);
}

f32 backward(network *net) {
  // Cross-entropy loss: L = -sum(target * log(out))
  f32 loss = 0.0f;
  for (u32 i = 0; i < OUTPUT_SIZE; i++) {
    if (net->target->data[i] > 0) {
      loss -= net->target->data[i] * logf(net->out->data[i] + 1e-7f);
    }
  }

  // dz2 = out - target (softmax + cross-entropy gradient)
  mat_sub(net->dz2, net->out, net->target);

  // dW2 = dz2 * a1^T, db2 = dz2
  mat_mul(net->dW2, net->dz2, net->a1, 0, 0, 1);
  mat_add(net->db2, net->db2, net->dz2);

  // da1 = W2^T * dz2
  mat_mul(net->da1, net->W2, net->dz2, 1, 1, 0);

  // dz1 = da1 * relu'(z1)
  for (u32 i = 0; i < HIDDEN_SIZE; i++) {
    net->dz1->data[i] = net->z1->data[i] > 0 ? net->da1->data[i] : 0;
  }

  // dW1 = dz1 * input^T, db1 = dz1
  mat_mul(net->dW1, net->dz1, net->input, 0, 0, 1);
  mat_add(net->db1, net->db1, net->dz1);

  return loss;
}

void update(network *net, f32 lr) {
  mat_scale(net->dW1, lr);
  mat_scale(net->db1, lr);
  mat_scale(net->dW2, lr);
  mat_scale(net->db2, lr);

  mat_sub(net->W1, net->W1, net->dW1);
  mat_sub(net->b1, net->b1, net->db1);
  mat_sub(net->W2, net->W2, net->dW2);
  mat_sub(net->b2, net->b2, net->db2);

  mat_clear(net->dW1);
  mat_clear(net->db1);
  mat_clear(net->dW2);
  mat_clear(net->db2);
}

// ============== Main ==============

int main(void) {
  printf("=== Text Sentiment Classification Test ===\n\n");

  mem_arena *arena = arena_create(MiB(64), MiB(1));
  network *net = create_network(arena);

  printf("Training on %zu samples...\n\n", NUM_SAMPLES);

  // Training
  u32 epochs = 500;
  f32 lr = 0.1f;

  for (u32 epoch = 0; epoch < epochs; epoch++) {
    f32 total_loss = 0.0f;
    u32 correct = 0;

    for (u32 i = 0; i < NUM_SAMPLES; i++) {
      // Encode text to vector
      encode_text(training_data[i].text, net->input->data);

      // Set target (one-hot)
      mat_clear(net->target);
      net->target->data[training_data[i].label] = 1.0f;

      // Forward pass
      forward(net);

      // Check prediction
      if (mat_argmax(net->out) == training_data[i].label) {
        correct++;
      }

      // Backward pass
      total_loss += backward(net);
    }

    // Update weights
    update(net, lr / NUM_SAMPLES);

    if ((epoch + 1) % 100 == 0) {
      printf("Epoch %3d: Loss = %.4f, Accuracy = %d/%zu (%.1f%%)\n", epoch + 1,
             total_loss / NUM_SAMPLES, correct, NUM_SAMPLES,
             (f32)correct / NUM_SAMPLES * 100.0f);
    }
  }

  printf("\n=== Testing on new sentences ===\n\n");

  // Test sentences
  const char *test_sentences[] = {
      "I love it",           "This is bad",  "Amazing quality",
      "Terrible experience", "Really great", "Very poor",
  };

  for (u32 i = 0; i < 6; i++) {
    encode_text(test_sentences[i], net->input->data);
    forward(net);

    u32 pred = mat_argmax(net->out);
    printf("\"%s\"\n", test_sentences[i]);
    printf("  → %s (%.1f%% positive, %.1f%% negative)\n\n",
           pred == 0 ? "POSITIVE" : "NEGATIVE", net->out->data[0] * 100.0f,
           net->out->data[1] * 100.0f);
  }

  arena_destroy(arena);
  return 0;
}