bump version in Cargo.toml

Refactor inference

dedicated crates for llama and gemma inferencing, not integrated

.github/dependabot.yml

@@ -0,0 +1,49 @@
+version: 2
+updates:
+  # Monitor Rust dependencies in the main crate
+  - package-ecosystem: "cargo"
+    directory: "/crates/predict-otron-9000"
+    schedule:
+      interval: "weekly"
+      day: "monday"
+      time: "09:00"
+      timezone: "UTC"
+    # Focus on security updates with higher priority
+    open-pull-requests-limit: 10
+    reviewers:
+      - "security-team"
+    assignees:
+      - "maintainer"
+    labels:
+      - "dependencies"
+      - "security"
+    # Security updates get higher priority
+    allow:
+      - dependency-type: "all"
+    # Group minor and patch updates to reduce noise
+    # Separate major updates for careful review
+    ignore:
+      - dependency-name: "*"
+        update-types: ["version-update:semver-major"]
+    commit-message:
+      prefix: "deps"
+      include: "scope"
+
+  # Monitor security updates more frequently
+  - package-ecosystem: "cargo"
+    directory: "/crates/predict-otron-9000"
+    schedule:
+      interval: "daily"
+    # Only security updates in daily checks
+    allow:
+      - dependency-type: "direct"
+        update-types: ["security"]
+      - dependency-type: "indirect"
+        update-types: ["security"]
+    open-pull-requests-limit: 5
+    labels:
+      - "security-update"
+      - "high-priority"
+    commit-message:
+      prefix: "security"
+      include: "scope"

.github/workflows/ci.yml

@@ -0,0 +1,47 @@
+name: CI
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  build:
+    name: build-and-test
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cargo/bin/
+            ~/.cargo/registry/index/
+            ~/.cargo/registry/cache/
+            ~/.cargo/git/db/
+            target/
+          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+
+      - name: Setup Rust
+        run: rustup update stable && rustup default stable
+
+      - name: Install clippy and rustfmt
+        run: rustup component add clippy rustfmt
+
+      - name: Cargo fmt (check)
+        run: cargo fmt --all -- --check
+
+      - name: Clippy
+        shell: bash
+        run: cargo clippy --all-targets
+
+      - name: Tests
+        shell: bash
+        run: cargo test --all
+
+      - name: Build Docs
+        shell: bash
+        run: | 
+          cargo doc -p predict-otron-9000 --no-deps 

.github/workflows/release.yml

@@ -0,0 +1,232 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  test:
+    name: Test before release
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: crates/predict-otron-9000
+    strategy:
+      fail-fast: false
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cargo/bin/
+            ~/.cargo/registry/index/
+            ~/.cargo/registry/cache/
+            ~/.cargo/git/db/
+            target/
+          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+
+      - name: Setup Rust
+        run: rustup update stable && rustup default stable
+
+      - name: Install clippy and rustfmt
+        run: rustup component add clippy rustfmt
+
+      - name: Cargo fmt (check)
+        run: cargo fmt --all -- --check
+
+      - name: Clippy
+        shell: bash
+        run: cargo clippy --all-targets
+
+      - name: Tests
+        shell: bash
+        run: cargo test --all
+
+#  publish:
+#    name: Publish to crates.io
+#    runs-on: ubuntu-latest
+#    permissions:
+#      id-token: write # Required for OIDC token exchange https://crates.io/docs/trusted-publishing
+#    needs: test
+#    defaults:
+#      run:
+#        working-directory: crates/predict-otron-9000
+#    steps:
+#      - name: Checkout
+#        uses: actions/checkout@v4
+#
+#      - uses: actions/cache@v4
+#        with:
+#          path: |
+#            ~/.cargo/bin/
+#            ~/.cargo/registry/index/
+#            ~/.cargo/registry/cache/
+#            ~/.cargo/git/db/
+#            target/
+#          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+#
+#      - name: Setup Rust
+#        run: rustup update stable && rustup default stable
+#
+#      - name: Verify tag matches version
+#        run: |
+#          TAG_VERSION=${GITHUB_REF#refs/tags/v}
+#          CARGO_VERSION=$(cargo metadata --no-deps --format-version 1 | jq -r '.packages[0].version')
+#          if [ "$TAG_VERSION" != "$CARGO_VERSION" ]; then
+#            echo "Tag version ($TAG_VERSION) does not match Cargo.toml version ($CARGO_VERSION)"
+#            exit 1
+#          fi
+#
+#      # See Trusted publishing: https://crates.io/docs/trusted-publishing
+#      - uses: rust-lang/crates-io-auth-action@v1
+#        id: auth
+#
+#      - run: cargo publish
+#        env:
+#          CARGO_REGISTRY_TOKEN: ${{ steps.auth.outputs.token }}
+
+  build-binaries:
+    name: Build binaries
+    runs-on: ${{ matrix.os }}
+    needs: test
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - target: x86_64-unknown-linux-gnu
+            os: ubuntu-latest
+            name: predict-otron-9000-x86_64-unknown-linux-gnu
+          - target: x86_64-apple-darwin
+            os: macos-latest
+            name: predict-otron-9000-x86_64-apple-darwin
+          - target: aarch64-apple-darwin
+            os: macos-latest
+            name: predict-otron-9000-aarch64-apple-darwin
+          - target: x86_64-pc-windows-msvc
+            os: windows-latest
+            name: predict-otron-9000-x86_64-pc-windows-msvc.exe
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cargo/bin/
+            ~/.cargo/registry/index/
+            ~/.cargo/registry/cache/
+            ~/.cargo/git/db/
+            target/
+          key: ${{ runner.os }}-${{ matrix.target }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+
+      - name: Setup Rust
+        run: rustup update stable && rustup default stable
+
+      - name: Add target
+        run: rustup target add ${{ matrix.target }}
+
+      - name: Build binary
+        run: cargo build --release --target ${{ matrix.target }} -p predict-otron-9000
+        env:
+          CARGO_TERM_COLOR: always
+
+      - name: Package binary (Unix)
+        if: matrix.os != 'windows-latest'
+        run: |
+          cd target/${{ matrix.target }}/release
+          tar czf ../../../${{ matrix.name }}.tar.gz predict-otron-9000
+          cd ../../../
+
+      - name: Package binary (Windows)
+        if: matrix.os == 'windows-latest'
+        run: |
+          cd target/${{ matrix.target }}/release
+          7z a ../../../${{ matrix.name }}.zip predict-otron-9000.exe
+          cd ../../../
+
+      - name: Upload binary artifacts (Unix)
+        if: matrix.os != 'windows-latest'
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.name }}
+          path: ${{ matrix.name }}.tar.gz
+
+      - name: Upload binary artifacts (Windows)
+        if: matrix.os == 'windows-latest'
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.name }}
+          path: ${{ matrix.name }}.zip
+
+  release:
+    name: Create GitHub Release
+    runs-on: ubuntu-latest
+    needs: [test, build-binaries]
+    permissions:
+      contents: write
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Extract tag name
+        id: tag
+        run: echo "tag=${GITHUB_REF#refs/tags/}" >> $GITHUB_OUTPUT
+
+      - name: Generate changelog
+        id: changelog
+        run: |
+          # Get the previous tag
+          PREV_TAG=$(git describe --tags --abbrev=0 HEAD^ 2>/dev/null || echo "")
+          
+          # Generate changelog
+          if [ -n "$PREV_TAG" ]; then
+            echo "## What's Changed" > changelog.md
+            echo "" >> changelog.md
+            git log --pretty=format:"* %s (%h)" ${PREV_TAG}..HEAD >> changelog.md
+            echo "" >> changelog.md
+            echo "" >> changelog.md
+            echo "**Full Changelog**: https://github.com/${{ github.repository }}/compare/${PREV_TAG}...${{ steps.tag.outputs.tag }}" >> changelog.md
+          else
+            echo "## What's Changed" > changelog.md
+            echo "" >> changelog.md
+            echo "Initial release of predict-otron-9000" >> changelog.md
+            echo "" >> changelog.md
+            echo "OpenAI Compatible Inference Server" >> changelog.md
+          fi
+          
+          # Set the changelog as output (handle multiline)
+          echo "changelog<<EOF" >> $GITHUB_OUTPUT
+          cat changelog.md >> $GITHUB_OUTPUT
+          echo "EOF" >> $GITHUB_OUTPUT
+
+      - name: Download all artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: artifacts
+
+      - name: Create Release
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          if [[ "${{ steps.tag.outputs.tag }}" == *"-"* ]]; then
+            PRERELEASE_FLAG="--prerelease"
+          else
+            PRERELEASE_FLAG=""
+          fi
+          
+          gh release create "${{ steps.tag.outputs.tag }}" \
+            --title "Release ${{ steps.tag.outputs.tag }}" \
+            --notes-file changelog.md \
+            $PRERELEASE_FLAG \
+            artifacts/predict-otron-9000-x86_64-unknown-linux-gnu/predict-otron-9000-x86_64-unknown-linux-gnu.tar.gz \
+            artifacts/predict-otron-9000-x86_64-apple-darwin/predict-otron-9000-x86_64-apple-darwin.tar.gz \
+            artifacts/predict-otron-9000-aarch64-apple-darwin/predict-otron-9000-aarch64-apple-darwin.tar.gz \
+            artifacts/predict-otron-9000-x86_64-pc-windows-msvc.exe/predict-otron-9000-x86_64-pc-windows-msvc.exe.zip

.gitignore

@@ -76,3 +76,5 @@ venv/
 *.bak
 *.backup
 *~
+/scripts/cli
+!/scripts/cli.ts

Cargo.toml

@@ -4,11 +4,41 @@ members = [
     "crates/inference-engine",
     "crates/embeddings-engine",
     "crates/leptos-app",
-    "crates/helm-chart-tool"
+    "crates/helm-chart-tool",
+    "crates/llama-runner",
+    "crates/gemma-runner"
 ]
 default-members = ["crates/predict-otron-9000"]
 resolver = "2"
 
+# Compiler optimization profiles for the workspace
+[profile.release]
+opt-level = 3
+debug = false
+strip = true
+lto = "thin"
+codegen-units = 1
+panic = "abort"
+
+[profile.dev]
+opt-level = 0
+debug = true
+strip = false
+overflow-checks = true
+
+# Profile for fast development builds with some optimization
+[profile.dev-opt]
+inherits = "dev"
+opt-level = 1
+debug = true
+overflow-checks = true
+
+# Profile for benchmarking and profiling
+[profile.bench]
+opt-level = 3
+debug = true
+lto = "thin"
+
 [[workspace.metadata.leptos]]
 # project name
 bin-package = "leptos-app"

README.md

@@ -1,16 +1,25 @@
-# predict-otron-9000
-
-A comprehensive multi-service AI platform built around local LLM inference, embeddings, and web interfaces.
-
+<h1 align="center">
+    predict-otron-9000
+</h1>
 <p align="center">
 Powerful local AI inference with OpenAI-compatible APIs
 </p>
 
+<br/>
+
+> This project is an educational aide for bootstrapping my understanding of language model inferencing at the lowest levels I can, serving as a "rubber-duck" solution for Kubernetes based performance-oriented inference capabilities on air-gapped networks.
+
+> By isolating application behaviors in components at the crate level, development reduces to a short feedback loop for validation and integration, ultimately smoothing the learning curve for scalable AI systems.
+Stability is currently best effort. Many models require unique configuration. When stability is achieved, this project will be promoted to the seemueller-io GitHub organization under a different name.
+
+A comprehensive multi-service AI platform built around local LLM inference, embeddings, and web interfaces.
+
+
 ## Project Overview
 
 The predict-otron-9000 is a flexible AI platform that provides:
 
-- **Local LLM Inference**: Run Gemma models locally with CPU or GPU acceleration
+- **Local LLM Inference**: Run Gemma and Llama models locally with CPU or GPU acceleration
 - **Embeddings Generation**: Create text embeddings with FastEmbed
 - **Web Interface**: Interact with models through a Leptos WASM chat interface
 - **TypeScript CLI**: Command-line client for testing and automation
@@ -22,31 +31,35 @@ The system supports both CPU and GPU acceleration (CUDA/Metal), with intelligent
 
 - **OpenAI Compatible**: API endpoints match OpenAI's format for easy integration
 - **Text Embeddings**: Generate high-quality text embeddings using FastEmbed
-- **Text Generation**: Chat completions with OpenAI-compatible API using Gemma models (1B, 2B, 7B variants including instruction-tuned models)
+- **Text Generation**: Chat completions with OpenAI-compatible API using Gemma and Llama models (various sizes including instruction-tuned variants)
 - **Performance Optimized**: Efficient caching and platform-specific optimizations for improved throughput
-- **Web Chat Interface**: Leptos-based WebAssembly (WASM) chat interface for browser-based interaction
+- **Web Chat Interface**: Leptos chat interface
 - **Flexible Deployment**: Run as monolithic service or microservices architecture
 
 ## Architecture Overview
 
 ### Workspace Structure
 
-The project uses a 4-crate Rust workspace plus TypeScript components:
+The project uses a 7-crate Rust workspace plus TypeScript components:
 
 ```
 crates/
 ├── predict-otron-9000/     # Main orchestration server (Rust 2024)
-├── inference-engine/       # Gemma inference via Candle (Rust 2021) 
-├── embeddings-engine/      # FastEmbed embeddings service (Rust 2024)
-└── leptos-app/             # WASM web frontend (Rust 2021)
-cli.ts                      # TypeScript/Bun CLI client
+├── inference-engine/       # Multi-model inference orchestrator (Rust 2021)
+├── gemma-runner/          # Gemma model inference via Candle (Rust 2021)
+├── llama-runner/          # Llama model inference via Candle (Rust 2021)
+├── embeddings-engine/     # FastEmbed embeddings service (Rust 2024)
+├── leptos-app/            # WASM web frontend (Rust 2021)
+├── helm-chart-tool/       # Kubernetes deployment tooling (Rust 2024)
+└── scripts/
+    └── cli.ts             # TypeScript/Bun CLI client
 ```
 
 ### Service Architecture
 
 - **Main Server** (port 8080): Orchestrates inference and embeddings services
 - **Embeddings Service** (port 8080): Standalone FastEmbed service with OpenAI API compatibility  
-- **Web Frontend** (port 8788): Leptos WASM chat interface served by Trunk
+- **Web Frontend** (port 8788): cargo leptos SSR app
 - **CLI Client**: TypeScript/Bun client for testing and automation
 
 ### Deployment Modes
@@ -149,16 +162,16 @@ cd crates/leptos-app
 #### TypeScript CLI Client
 ```bash
 # List available models
-bun run cli.ts --list-models
+bun run scripts/cli.ts --list-models
 
 # Chat completion
-bun run cli.ts "What is the capital of France?"
+bun run scripts/cli.ts "What is the capital of France?"
 
 # With specific model
-bun run cli.ts --model gemma-3-1b-it --prompt "Hello, world!"
+bun run scripts/cli.ts --model gemma-3-1b-it --prompt "Hello, world!"
 
 # Show help
-bun run cli.ts --help
+bun run scripts/cli.ts --help
 ```
 
 ## API Usage
@@ -274,7 +287,7 @@ cargo test --workspace
 
 **End-to-end test script:**
 ```bash
-./test.sh
+./smoke_test.sh
 ```
 
 This script:
@@ -454,7 +467,7 @@ curl -s http://localhost:8080/v1/models | jq
 
 **CLI client test:**
 ```bash
-bun run cli.ts "What is 2+2?"
+bun run scripts/cli.ts "What is 2+2?"
 ```
 
 **Web frontend:**
@@ -465,7 +478,7 @@ cd crates/leptos-app && ./run.sh &
 
 **Integration test:**
 ```bash
-./test.sh
+./smoke_test.sh
 ```
 
 **Cleanup:**
@@ -493,4 +506,4 @@ For networked tests and full functionality, ensure Hugging Face authentication i
 4. Ensure all tests pass: `cargo test`
 5. Submit a pull request
 
-_Warning: Do NOT use this in production unless you are cool like that._
+_Warning: Do NOT use this in production unless you are cool like that._

crates/embeddings-engine/src/lib.rs

@@ -1,9 +1,5 @@
 use async_openai::types::{CreateEmbeddingRequest, EmbeddingInput};
-use axum::{
-    response::Json as ResponseJson, routing::{post},
-    Json,
-    Router,
-};
+use axum::{Json, Router, response::Json as ResponseJson, routing::post};
 use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
 use once_cell::sync::Lazy;
 use tower_http::trace::TraceLayer;
@@ -13,15 +9,18 @@ use tracing;
 static EMBEDDING_MODEL: Lazy<TextEmbedding> = Lazy::new(|| {
     tracing::info!("Initializing persistent embedding model (singleton)");
     let model_start_time = std::time::Instant::now();
-    
+
     let model = TextEmbedding::try_new(
-        InitOptions::new(EmbeddingModel::NomicEmbedTextV15).with_show_download_progress(true)
+        InitOptions::new(EmbeddingModel::NomicEmbedTextV15).with_show_download_progress(true),
     )
-        .expect("Failed to initialize persistent embedding model");
-    
+    .expect("Failed to initialize persistent embedding model");
+
     let model_init_time = model_start_time.elapsed();
-    tracing::info!("Persistent embedding model initialized in {:.2?}", model_init_time);
-    
+    tracing::info!(
+        "Persistent embedding model initialized in {:.2?}",
+        model_init_time
+    );
+
     model
 });
 
@@ -30,18 +29,21 @@ pub async fn embeddings_create(
 ) -> ResponseJson<serde_json::Value> {
     // Start timing the entire process
     let start_time = std::time::Instant::now();
-    
+
     // Phase 1: Access persistent model instance
     let model_start_time = std::time::Instant::now();
-    
+
     // Access the lazy-initialized persistent model instance
     // This will only initialize the model on the first request
     let model_access_time = model_start_time.elapsed();
-    tracing::debug!("Persistent model access completed in {:.2?}", model_access_time);
-    
+    tracing::debug!(
+        "Persistent model access completed in {:.2?}",
+        model_access_time
+    );
+
     // Phase 2: Process input
     let input_start_time = std::time::Instant::now();
-    
+
     let embedding_input = payload.input;
     let texts_from_embedding_input = match embedding_input {
         EmbeddingInput::String(text) => vec![text],
@@ -53,41 +55,58 @@ pub async fn embeddings_create(
             panic!("Array of integer arrays not supported for text embeddings");
         }
     };
-    
+
     let input_processing_time = input_start_time.elapsed();
-    tracing::debug!("Input processing completed in {:.2?}", input_processing_time);
-    
+    tracing::debug!(
+        "Input processing completed in {:.2?}",
+        input_processing_time
+    );
+
     // Phase 3: Generate embeddings
     let embedding_start_time = std::time::Instant::now();
-    
+
     let embeddings = EMBEDDING_MODEL
         .embed(texts_from_embedding_input, None)
         .expect("failed to embed document");
-    
+
     let embedding_generation_time = embedding_start_time.elapsed();
-    tracing::info!("Embedding generation completed in {:.2?}", embedding_generation_time);
-    
+    tracing::info!(
+        "Embedding generation completed in {:.2?}",
+        embedding_generation_time
+    );
+
     // Memory usage estimation (approximate)
-    let embedding_size_bytes = embeddings.iter()
+    let embedding_size_bytes = embeddings
+        .iter()
         .map(|e| e.len() * std::mem::size_of::<f32>())
         .sum::<usize>();
-    tracing::debug!("Embedding size: {:.2} MB", embedding_size_bytes as f64 / 1024.0 / 1024.0);
+    tracing::debug!(
+        "Embedding size: {:.2} MB",
+        embedding_size_bytes as f64 / 1024.0 / 1024.0
+    );
 
     // Only log detailed embedding information at trace level to reduce log volume
     tracing::trace!("Embeddings length: {}", embeddings.len());
     tracing::info!("Embedding dimension: {}", embeddings[0].len());
 
     // Log the first 10 values of the original embedding at trace level
-    tracing::trace!("Original embedding preview: {:?}", &embeddings[0][..10.min(embeddings[0].len())]);
+    tracing::trace!(
+        "Original embedding preview: {:?}",
+        &embeddings[0][..10.min(embeddings[0].len())]
+    );
 
     // Check if there are any NaN or zero values in the original embedding
     let nan_count = embeddings[0].iter().filter(|&&x| x.is_nan()).count();
     let zero_count = embeddings[0].iter().filter(|&&x| x == 0.0).count();
-    tracing::trace!("Original embedding stats: NaN count={}, zero count={}", nan_count, zero_count);
+    tracing::trace!(
+        "Original embedding stats: NaN count={}, zero count={}",
+        nan_count,
+        zero_count
+    );
 
     // Phase 4: Post-process embeddings
     let postprocessing_start_time = std::time::Instant::now();
-    
+
     // Create the final embedding
     let final_embedding = {
         // Check if the embedding is all zeros
@@ -110,6 +129,8 @@ pub async fn embeddings_create(
 
             // Normalize the random embedding
             let norm: f32 = random_embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
+
+            #[allow(clippy::needless_range_loop)]
             for i in 0..random_embedding.len() {
                 random_embedding[i] /= norm;
             }
@@ -123,25 +144,35 @@ pub async fn embeddings_create(
             let target_dimension = 768;
             if padded_embedding.len() < target_dimension {
                 let padding_needed = target_dimension - padded_embedding.len();
-                tracing::trace!("Padding embedding with {} zeros to reach {} dimensions", padding_needed, target_dimension);
+                tracing::trace!(
+                    "Padding embedding with {} zeros to reach {} dimensions",
+                    padding_needed,
+                    target_dimension
+                );
                 padded_embedding.extend(vec![0.0; padding_needed]);
             }
 
             padded_embedding
         }
     };
-    
+
     let postprocessing_time = postprocessing_start_time.elapsed();
-    tracing::debug!("Embedding post-processing completed in {:.2?}", postprocessing_time);
+    tracing::debug!(
+        "Embedding post-processing completed in {:.2?}",
+        postprocessing_time
+    );
 
     tracing::trace!("Final embedding dimension: {}", final_embedding.len());
 
     // Log the first 10 values of the final embedding at trace level
-    tracing::trace!("Final embedding preview: {:?}", &final_embedding[..10.min(final_embedding.len())]);
+    tracing::trace!(
+        "Final embedding preview: {:?}",
+        &final_embedding[..10.min(final_embedding.len())]
+    );
 
     // Phase 5: Prepare response
     let response_start_time = std::time::Instant::now();
-    
+
     // Return a response that matches the OpenAI API format
     let response = serde_json::json!({
         "object": "list",
@@ -158,10 +189,10 @@ pub async fn embeddings_create(
             "total_tokens": 0
         }
     });
-    
+
     let response_time = response_start_time.elapsed();
     tracing::debug!("Response preparation completed in {:.2?}", response_time);
-    
+
     // Log total time and breakdown
     let total_time = start_time.elapsed();
     tracing::info!(
@@ -171,7 +202,7 @@ pub async fn embeddings_create(
         embedding_generation_time,
         postprocessing_time
     );
-    
+
     ResponseJson(response)
 }
 
@@ -179,4 +210,4 @@ pub fn create_embeddings_router() -> Router {
     Router::new()
         .route("/v1/embeddings", post(embeddings_create))
         .layer(TraceLayer::new_for_http())
-}
+}

crates/embeddings-engine/src/main.rs

@@ -1,8 +1,8 @@
 use async_openai::types::{CreateEmbeddingRequest, EmbeddingInput};
 use axum::{
-	response::Json as ResponseJson, routing::{get, post},
-	Json,
-	Router,
+    Json, Router,
+    response::Json as ResponseJson,
+    routing::{get, post},
 };
 use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
 use serde::{Deserialize, Serialize};
@@ -13,19 +13,17 @@ use tracing;
 const DEFAULT_SERVER_HOST: &str = "127.0.0.1";
 const DEFAULT_SERVER_PORT: &str = "8080";
 
-
 async fn embeddings_create(
     Json(payload): Json<CreateEmbeddingRequest>,
 ) -> ResponseJson<serde_json::Value> {
     let model = TextEmbedding::try_new(
-        InitOptions::new(EmbeddingModel::NomicEmbedTextV15).with_show_download_progress(true)
+        InitOptions::new(EmbeddingModel::NomicEmbedTextV15).with_show_download_progress(true),
     )
     .expect("Failed to initialize model");
 
+    let embedding_input = payload.input;
 
-	let embedding_input = payload.input;
-
-	let texts_from_embedding_input = match embedding_input {
+    let texts_from_embedding_input = match embedding_input {
         EmbeddingInput::String(text) => vec![text],
         EmbeddingInput::StringArray(texts) => texts,
         EmbeddingInput::IntegerArray(_) => {
@@ -45,12 +43,19 @@ async fn embeddings_create(
     tracing::info!("Embedding dimension: {}", embeddings[0].len());
 
     // Log the first 10 values of the original embedding at trace level
-    tracing::trace!("Original embedding preview: {:?}", &embeddings[0][..10.min(embeddings[0].len())]);
+    tracing::trace!(
+        "Original embedding preview: {:?}",
+        &embeddings[0][..10.min(embeddings[0].len())]
+    );
 
     // Check if there are any NaN or zero values in the original embedding
     let nan_count = embeddings[0].iter().filter(|&&x| x.is_nan()).count();
     let zero_count = embeddings[0].iter().filter(|&&x| x == 0.0).count();
-    tracing::trace!("Original embedding stats: NaN count={}, zero count={}", nan_count, zero_count);
+    tracing::trace!(
+        "Original embedding stats: NaN count={}, zero count={}",
+        nan_count,
+        zero_count
+    );
 
     // Create the final embedding
     let final_embedding = {
@@ -87,7 +92,11 @@ async fn embeddings_create(
             let target_dimension = 768;
             if padded_embedding.len() < target_dimension {
                 let padding_needed = target_dimension - padded_embedding.len();
-                tracing::trace!("Padding embedding with {} zeros to reach {} dimensions", padding_needed, target_dimension);
+                tracing::trace!(
+                    "Padding embedding with {} zeros to reach {} dimensions",
+                    padding_needed,
+                    target_dimension
+                );
                 padded_embedding.extend(vec![0.0; padding_needed]);
             }
 
@@ -98,7 +107,10 @@ async fn embeddings_create(
     tracing::trace!("Final embedding dimension: {}", final_embedding.len());
 
     // Log the first 10 values of the final embedding at trace level
-    tracing::trace!("Final embedding preview: {:?}", &final_embedding[..10.min(final_embedding.len())]);
+    tracing::trace!(
+        "Final embedding preview: {:?}",
+        &final_embedding[..10.min(final_embedding.len())]
+    );
 
     // Return a response that matches the OpenAI API format
     let response = serde_json::json!({
@@ -120,7 +132,7 @@ async fn embeddings_create(
 }
 
 fn create_app() -> Router {
-	Router::new()
+    Router::new()
         .route("/v1/embeddings", post(embeddings_create))
         .layer(TraceLayer::new_for_http())
 }
@@ -143,21 +155,21 @@ async fn main() {
         .init();
     let app = create_app();
 
-	let server_host = env::var("SERVER_HOST").unwrap_or_else(|_| DEFAULT_SERVER_HOST.to_string());
-	let server_port = env::var("SERVER_PORT").unwrap_or_else(|_| DEFAULT_SERVER_PORT.to_string());
-	let server_address = format!("{}:{}", server_host, server_port);
-	let listener = tokio::net::TcpListener::bind(server_address).await.unwrap();
-	tracing::info!("Listening on {}", listener.local_addr().unwrap());
+    let server_host = env::var("SERVER_HOST").unwrap_or_else(|_| DEFAULT_SERVER_HOST.to_string());
+    let server_port = env::var("SERVER_PORT").unwrap_or_else(|_| DEFAULT_SERVER_PORT.to_string());
+    let server_address = format!("{}:{}", server_host, server_port);
+    let listener = tokio::net::TcpListener::bind(server_address).await.unwrap();
+    tracing::info!("Listening on {}", listener.local_addr().unwrap());
     axum::serve(listener, app).await.unwrap();
 }
 
 #[cfg(test)]
 mod tests {
-	use super::*;
-	use axum::body::to_bytes;
-	use axum::body::Body;
-	use axum::http::StatusCode;
-	use tower::ServiceExt;
+    use super::*;
+    use axum::body::Body;
+    use axum::body::to_bytes;
+    use axum::http::StatusCode;
+    use tower::ServiceExt;
 
     #[tokio::test]
     async fn test_embeddings_create() {
@@ -168,11 +180,13 @@ mod tests {
 
         let body = CreateEmbeddingRequest {
             model: "nomic-text-embed".to_string(),
-            input: EmbeddingInput::from(vec!["The food was delicious and the waiter...".to_string()]),
-			encoding_format: None,
-			user: None,
-			dimensions: Some(768),
-		};
+            input: EmbeddingInput::from(vec![
+                "The food was delicious and the waiter...".to_string(),
+            ]),
+            encoding_format: None,
+            user: None,
+            dimensions: Some(768),
+        };
 
         let response = app
             .oneshot(

crates/gemma-runner/Cargo.toml

@@ -0,0 +1,32 @@
+[package]
+name = "gemma-runner"
+version = "0.1.0"
+edition = "2021"
+
+
+
+
+[dependencies]
+candle-core = { git = "https://github.com/huggingface/candle.git" }
+candle-nn = { git = "https://github.com/huggingface/candle.git" }
+candle-transformers = { git = "https://github.com/huggingface/candle.git" }
+candle-examples = { git = "https://github.com/huggingface/candle.git" }
+hf-hub = "0.4"
+tokenizers = "0.21"
+anyhow = "1.0"
+clap = { version = "4.0", features = ["derive", "string"] }
+serde_json = "1.0"
+tracing = "0.1"
+tracing-chrome = "0.7"
+tracing-subscriber = "0.3"
+
+[target.'cfg(target_os = "macos")'.dependencies]
+candle-core = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+candle-nn = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+candle-transformers = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+
+
+[features]
+default = []
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+metal = ["candle-core/metal", "candle-nn/metal", "candle-transformers/metal"]

crates/gemma-runner/README.md

@@ -0,0 +1,137 @@
+# Gemma Runner
+
+Fast Gemma inference with Candle framework in Rust.
+
+## Features
+
+- Support for multiple Gemma model versions (v1, v2, v3)
+- GPU acceleration with CUDA and Metal
+- Configurable sampling parameters
+- Multiple model variants including instruct and code models
+
+## Supported Models
+
+### Gemma v1
+- `gemma-2b` - Base 2B model
+- `gemma-7b` - Base 7B model  
+- `gemma-2b-it` - Instruct 2B model
+- `gemma-7b-it` - Instruct 7B model
+- `gemma-1.1-2b-it` - Instruct 2B v1.1 model
+- `gemma-1.1-7b-it` - Instruct 7B v1.1 model
+
+### CodeGemma
+- `codegemma-2b` - Code base 2B model
+- `codegemma-7b` - Code base 7B model
+- `codegemma-2b-it` - Code instruct 2B model
+- `codegemma-7b-it` - Code instruct 7B model
+
+### Gemma v2
+- `gemma-2-2b` - Base 2B v2 model (default)
+- `gemma-2-2b-it` - Instruct 2B v2 model
+- `gemma-2-9b` - Base 9B v2 model
+- `gemma-2-9b-it` - Instruct 9B v2 model
+
+### Gemma v3
+- `gemma-3-1b` - Base 1B v3 model
+- `gemma-3-1b-it` - Instruct 1B v3 model
+
+## Installation
+
+```bash
+cd gemma-runner
+cargo build --release
+```
+
+For GPU support:
+```bash
+# CUDA
+cargo build --release --features cuda
+
+# Metal (macOS)
+cargo build --release --features metal
+```
+
+## Usage
+
+### Basic Usage
+
+```bash
+# Run with default model (gemma-2-2b)
+cargo run -- --prompt "The capital of France is"
+
+# Specify a different model
+cargo run -- --model gemma-2b-it --prompt "Explain quantum computing"
+
+# Generate more tokens
+cargo run -- --model codegemma-2b-it --prompt "Write a Python function to sort a list" --max-tokens 200
+```
+
+### Advanced Options
+
+```bash
+# Use CPU instead of GPU
+cargo run -- --cpu --prompt "Hello world"
+
+# Adjust sampling parameters
+cargo run -- --temperature 0.8 --top-p 0.9 --prompt "Write a story about"
+
+# Use custom model from HuggingFace Hub
+cargo run -- --model-id "google/gemma-2-2b-it" --prompt "What is AI?"
+
+# Enable tracing for performance analysis
+cargo run -- --tracing --prompt "Explain machine learning"
+```
+
+### Command Line Arguments
+
+- `--prompt, -p` - The prompt to generate text from (default: "The capital of France is")
+- `--model, -m` - The model to use (default: "gemma-2-2b")
+- `--cpu` - Run on CPU rather than GPU
+- `--temperature, -t` - Sampling temperature (optional)
+- `--top-p` - Nucleus sampling probability cutoff (optional)
+- `--seed` - Random seed (default: 299792458)
+- `--max-tokens, -n` - Maximum tokens to generate (default: 100)
+- `--model-id` - Custom model ID from HuggingFace Hub
+- `--revision` - Model revision (default: "main")
+- `--use-flash-attn` - Use flash attention
+- `--repeat-penalty` - Repetition penalty (default: 1.1)
+- `--repeat-last-n` - Context size for repeat penalty (default: 64)
+- `--dtype` - Data type (f16, bf16, f32)
+- `--tracing` - Enable performance tracing
+
+## Examples
+
+### Text Generation
+```bash
+cargo run -- --model gemma-2b-it --prompt "Explain the theory of relativity" --max-tokens 150
+```
+
+### Code Generation
+```bash
+cargo run -- --model codegemma-7b-it --prompt "Write a Rust function to calculate factorial" --max-tokens 100
+```
+
+### Creative Writing
+```bash
+cargo run -- --model gemma-7b-it --temperature 0.9 --prompt "Once upon a time in a magical forest" --max-tokens 200
+```
+
+### Chat with Gemma 3 (Instruct format)
+```bash
+cargo run -- --model gemma-3-1b-it --prompt "How do I learn Rust programming?"
+```
+
+## Performance Notes
+
+- GPU acceleration is automatically detected and used when available
+- BF16 precision is used on CUDA for better performance
+- F32 precision is used on CPU
+- Flash attention can be enabled with `--use-flash-attn` for supported models
+- Model files are cached locally after first download
+
+## Requirements
+
+- Rust 1.70+
+- CUDA toolkit (for CUDA support)
+- Metal (automatically available on macOS)
+- Internet connection for first-time model download

crates/gemma-runner/src/gemma_api.rs

@@ -0,0 +1,397 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+use anyhow::{Error as E, Result};
+use candle_transformers::models::gemma::{Config as Config1, Model as Model1};
+use candle_transformers::models::gemma2::{Config as Config2, Model as Model2};
+use candle_transformers::models::gemma3::{Config as Config3, Model as Model3};
+use clap::ValueEnum;
+
+// Removed gemma_cli import as it's not needed for the API
+use candle_core::{utils, DType, Device, Tensor};
+use candle_examples::token_output_stream::TokenOutputStream;
+use candle_nn::VarBuilder;
+use candle_transformers::generation::LogitsProcessor;
+use hf_hub::{api::sync::Api, Repo, RepoType};
+use std::io::Write;
+use tokenizers::Tokenizer;
+
+use std::sync::mpsc::{self, Receiver, Sender};
+use std::thread;
+
+#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)]
+pub enum WhichModel {
+    #[value(name = "gemma-2b")]
+    Base2B,
+    #[value(name = "gemma-7b")]
+    Base7B,
+    #[value(name = "gemma-2b-it")]
+    Instruct2B,
+    #[value(name = "gemma-7b-it")]
+    Instruct7B,
+    #[value(name = "gemma-1.1-2b-it")]
+    InstructV1_1_2B,
+    #[value(name = "gemma-1.1-7b-it")]
+    InstructV1_1_7B,
+    #[value(name = "codegemma-2b")]
+    CodeBase2B,
+    #[value(name = "codegemma-7b")]
+    CodeBase7B,
+    #[value(name = "codegemma-2b-it")]
+    CodeInstruct2B,
+    #[value(name = "codegemma-7b-it")]
+    CodeInstruct7B,
+    #[value(name = "gemma-2-2b")]
+    BaseV2_2B,
+    #[value(name = "gemma-2-2b-it")]
+    InstructV2_2B,
+    #[value(name = "gemma-2-9b")]
+    BaseV2_9B,
+    #[value(name = "gemma-2-9b-it")]
+    InstructV2_9B,
+    #[value(name = "gemma-3-1b")]
+    BaseV3_1B,
+    #[value(name = "gemma-3-1b-it")]
+    InstructV3_1B,
+}
+
+enum Model {
+    V1(Model1),
+    V2(Model2),
+    V3(Model3),
+}
+
+impl Model {
+    fn forward(&mut self, input_ids: &Tensor, pos: usize) -> candle_core::Result<Tensor> {
+        match self {
+            Self::V1(m) => m.forward(input_ids, pos),
+            Self::V2(m) => m.forward(input_ids, pos),
+            Self::V3(m) => m.forward(input_ids, pos),
+        }
+    }
+}
+
+pub struct TextGeneration {
+    model: Model,
+    device: Device,
+    tokenizer: TokenOutputStream,
+    logits_processor: LogitsProcessor,
+    repeat_penalty: f32,
+    repeat_last_n: usize,
+}
+
+fn device(cpu: bool) -> Result<Device> {
+    if cpu {
+        Ok(Device::Cpu)
+    } else if utils::cuda_is_available() {
+        Ok(Device::new_cuda(0)?)
+    } else if utils::metal_is_available() {
+        Ok(Device::new_metal(0)?)
+    } else {
+        Ok(Device::Cpu)
+    }
+}
+
+impl TextGeneration {
+    #[allow(clippy::too_many_arguments)]
+    fn new(
+        model: Model,
+        tokenizer: Tokenizer,
+        seed: u64,
+        temp: Option<f64>,
+        top_p: Option<f64>,
+        repeat_penalty: f32,
+        repeat_last_n: usize,
+        device: &Device,
+    ) -> Self {
+        let logits_processor = LogitsProcessor::new(seed, temp, top_p);
+        Self {
+            model,
+            tokenizer: TokenOutputStream::new(tokenizer),
+            logits_processor,
+            repeat_penalty,
+            repeat_last_n,
+            device: device.clone(),
+        }
+    }
+
+    /// Stream-only generation: sends freshly generated token strings over `tx`.
+    /// (Does not send the prompt tokens; only newly generated model tokens.)
+    fn run_stream(
+        &mut self,
+        prompt: &str,
+        sample_len: usize,
+        tx: Sender<Result<String>>,
+    ) -> Result<()> {
+        self.tokenizer.clear();
+
+        // Encode prompt (context only; do not emit prompt tokens to the stream).
+        let mut tokens = self
+            .tokenizer
+            .tokenizer()
+            .encode(prompt, true)
+            .map_err(E::msg)?
+            .get_ids()
+            .to_vec();
+
+        // Warm the tokenizer's internal state with prompt tokens (so merges are correct),
+        // but do not send them to the receiver.
+        for &t in tokens.iter() {
+            let _ = self.tokenizer.next_token(t)?;
+        }
+        // Make sure stdout isn't holding anything (if caller also prints).
+        std::io::stdout().flush()?;
+
+        let mut generated_tokens = 0usize;
+
+        let eos_token = match self.tokenizer.get_token("<eos>") {
+            Some(token) => token,
+            None => anyhow::bail!("cannot find the <eos> token"),
+        };
+        let eot_token = match self.tokenizer.get_token("<end_of_turn>") {
+            Some(token) => token,
+            None => {
+                eprintln!("Warning: <end_of_turn> token not found, using <eos> as backup");
+                eos_token
+            }
+        };
+
+        let start_gen = std::time::Instant::now();
+
+        for index in 0..sample_len {
+            let context_size = if index > 0 { 1 } else { tokens.len() };
+            let start_pos = tokens.len().saturating_sub(context_size);
+            let ctxt = &tokens[start_pos..];
+
+            let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
+            let logits = self.model.forward(&input, start_pos)?;
+            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+
+            let logits = if self.repeat_penalty == 1. {
+                logits
+            } else {
+                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
+                candle_transformers::utils::apply_repeat_penalty(
+                    &logits,
+                    self.repeat_penalty,
+                    &tokens[start_at..],
+                )?
+            };
+
+            let next_token = self.logits_processor.sample(&logits)?;
+            tokens.push(next_token);
+            generated_tokens += 1;
+
+            if next_token == eos_token || next_token == eot_token {
+                break;
+            }
+
+            if let Some(t) = self.tokenizer.next_token(next_token)? {
+                // Best-effort send; ignore if receiver dropped.
+                let _ = tx.send(Ok(t));
+            }
+        }
+
+        let _dt = start_gen.elapsed();
+
+        // Flush any remaining buffered bytes as one final chunk.
+        if let Some(rest) = self.tokenizer.decode_rest().map_err(E::msg)? {
+            let _ = tx.send(Ok(rest));
+        }
+
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct GemmaInferenceConfig {
+    pub tracing: bool,
+    pub prompt: String,
+    pub model: WhichModel,
+    pub cpu: bool,
+    pub dtype: Option<String>,
+    pub model_id: Option<String>,
+    pub revision: String,
+    pub use_flash_attn: bool,
+    pub seed: u64,
+    pub temperature: f64,
+    pub top_p: Option<f64>,
+    pub repeat_penalty: f32,
+    pub repeat_last_n: usize,
+    pub max_tokens: usize,
+}
+
+impl Default for GemmaInferenceConfig {
+    fn default() -> Self {
+        Self {
+            tracing: false,
+            prompt: "Hello".to_string(),
+            model: WhichModel::InstructV2_2B,
+            cpu: false,
+            dtype: None,
+            model_id: None,
+            revision: "main".to_string(),
+            use_flash_attn: false,
+            seed: 299792458,
+            temperature: 0.8,
+            top_p: None,
+            repeat_penalty: 1.1,
+            repeat_last_n: 128,
+            max_tokens: 100,
+        }
+    }
+}
+
+// Removed From<Args> implementation as Args is not available and not needed for API usage
+
+/// Builds the model and returns a channel that streams generated token strings.
+/// If model setup fails, the `Result` is returned immediately.
+pub fn run_gemma_api(cfg: GemmaInferenceConfig) -> Result<Receiver<Result<String>>> {
+    use tracing_chrome::ChromeLayerBuilder;
+    use tracing_subscriber::prelude::*;
+
+    let _guard = if cfg.tracing {
+        let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
+        tracing_subscriber::registry().with(chrome_layer).init();
+        Some(guard)
+    } else {
+        None
+    };
+
+    println!(
+        "avx: {}, neon: {}, simd128: {}, f16c: {}",
+        utils::with_avx(),
+        utils::with_neon(),
+        utils::with_simd128(),
+        utils::with_f16c()
+    );
+
+    let device = device(cfg.cpu)?;
+    println!("Device: {:?}", device);
+
+    let dtype = match cfg.dtype.as_deref() {
+        Some("f16") => DType::F16,
+        Some("bf16") => DType::BF16,
+        Some("f32") => DType::F32,
+        Some(dtype) => anyhow::bail!("Unsupported dtype {dtype}"),
+        None => {
+            if device.is_cuda() {
+                DType::BF16
+            } else {
+                DType::F16
+            }
+        }
+    };
+    println!("Using dtype: {:?}", dtype);
+
+    let start = std::time::Instant::now();
+    let api = Api::new()?;
+
+    let model_id = cfg.model_id.unwrap_or_else(|| {
+        match cfg.model {
+            WhichModel::Base2B => "google/gemma-2b",
+            WhichModel::Base7B => "google/gemma-7b",
+            WhichModel::Instruct2B => "google/gemma-2b-it",
+            WhichModel::Instruct7B => "google/gemma-7b-it",
+            WhichModel::InstructV1_1_2B => "google/gemma-1.1-2b-it",
+            WhichModel::InstructV1_1_7B => "google/gemma-1.1-7b-it",
+            WhichModel::CodeBase2B => "google/codegemma-2b",
+            WhichModel::CodeBase7B => "google/codegemma-7b",
+            WhichModel::CodeInstruct2B => "google/codegemma-2b-it",
+            WhichModel::CodeInstruct7B => "google/codegemma-7b-it",
+            WhichModel::BaseV2_2B => "google/gemma-2-2b",
+            WhichModel::InstructV2_2B => "google/gemma-2-2b-it",
+            WhichModel::BaseV2_9B => "google/gemma-2-9b",
+            WhichModel::InstructV2_9B => "google/gemma-2-9b-it",
+            WhichModel::BaseV3_1B => "google/gemma-3-1b-pt",
+            WhichModel::InstructV3_1B => "google/gemma-3-1b-it",
+        }
+        .to_string()
+    });
+
+    println!("Loading model: {}", &model_id);
+
+    let repo = api.repo(Repo::with_revision(model_id, RepoType::Model, cfg.revision));
+    let tokenizer_filename = repo.get("tokenizer.json")?;
+    let config_filename = repo.get("config.json")?;
+    let filenames = match cfg.model {
+        WhichModel::BaseV3_1B | WhichModel::InstructV3_1B => vec![repo.get("model.safetensors")?],
+        _ => candle_examples::hub_load_safetensors(&repo, "model.safetensors.index.json")?,
+    };
+    println!("Retrieved files in {:?}", start.elapsed());
+
+    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+
+    let start = std::time::Instant::now();
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, dtype, &device)? };
+
+    let model: Model = match cfg.model {
+        WhichModel::Base2B
+        | WhichModel::Base7B
+        | WhichModel::Instruct2B
+        | WhichModel::Instruct7B
+        | WhichModel::InstructV1_1_2B
+        | WhichModel::InstructV1_1_7B
+        | WhichModel::CodeBase2B
+        | WhichModel::CodeBase7B
+        | WhichModel::CodeInstruct2B
+        | WhichModel::CodeInstruct7B => {
+            let config: Config1 = serde_json::from_reader(std::fs::File::open(config_filename)?)?;
+            let model = Model1::new(cfg.use_flash_attn, &config, vb)?;
+            Model::V1(model)
+        }
+        WhichModel::BaseV2_2B
+        | WhichModel::InstructV2_2B
+        | WhichModel::BaseV2_9B
+        | WhichModel::InstructV2_9B => {
+            let config: Config2 = serde_json::from_reader(std::fs::File::open(config_filename)?)?;
+            let model = Model2::new(cfg.use_flash_attn, &config, vb)?;
+            Model::V2(model)
+        }
+        WhichModel::BaseV3_1B | WhichModel::InstructV3_1B => {
+            let config: Config3 = serde_json::from_reader(std::fs::File::open(config_filename)?)?;
+            let model = Model3::new(cfg.use_flash_attn, &config, vb)?;
+            Model::V3(model)
+        }
+    };
+    println!("Loaded model in {:?}", start.elapsed());
+
+    let mut pipeline = TextGeneration::new(
+        model,
+        tokenizer,
+        cfg.seed,
+        cfg.temperature.into(),
+        cfg.top_p,
+        cfg.repeat_penalty,
+        cfg.repeat_last_n,
+        &device,
+    );
+
+    let prompt = match cfg.model {
+        WhichModel::InstructV3_1B => {
+            format!(
+                "<start_of_turn>user\n{}<end_of_turn>\n<start_of_turn>model\n",
+                cfg.prompt
+            )
+        }
+        _ => cfg.prompt,
+    };
+
+    println!("Starting inference...");
+
+    // Create the channel after successful setup.
+    let (tx, rx) = mpsc::channel::<Result<String>>();
+
+    // Spawn generation thread; send tokens to the channel.
+    thread::spawn(move || {
+        // If generation fails, forward the error once.
+        if let Err(e) = pipeline.run_stream(&prompt, cfg.max_tokens, tx.clone()) {
+            let _ = tx.send(Err(e));
+        }
+        // Channel closes when tx is dropped.
+    });
+
+    Ok(rx)
+}

crates/gemma-runner/src/gemma_cli.rs

@@ -0,0 +1,97 @@
+use crate::gemma_api::{run_gemma_api, GemmaInferenceConfig, WhichModel};
+use clap::Parser;
+use std::io::Write;
+
+#[derive(Parser, Debug)]
+#[command(author, version, about = "Fast Gemma inference with Candle", long_about = None)]
+pub struct Args {
+    /// The prompt to generate text from
+    #[arg(short, long, default_value = "The capital of France is")]
+    pub(crate) prompt: String,
+
+    /// The model to use
+    #[arg(short, long, default_value = "gemma-2-2b")]
+    pub(crate) model: WhichModel,
+
+    /// Run on CPU rather than GPU
+    #[arg(long)]
+    pub(crate) cpu: bool,
+
+    /// The temperature used to generate samples
+    #[arg(short, long)]
+    pub(crate) temperature: Option<f64>,
+
+    /// Nucleus sampling probability cutoff
+    #[arg(long)]
+    pub(crate) top_p: Option<f64>,
+
+    /// The seed to use when generating random samples
+    #[arg(long, default_value_t = 299792458)]
+    pub(crate) seed: u64,
+
+    /// The length of the sample to generate (in tokens)
+    #[arg(short = 'n', long, default_value_t = 100)]
+    pub(crate) max_tokens: usize,
+
+    /// Use different dtype than default
+    #[arg(long)]
+    pub(crate) dtype: Option<String>,
+
+    /// Custom model ID from HuggingFace Hub
+    #[arg(long)]
+    pub(crate) model_id: Option<String>,
+
+    /// Model revision
+    #[arg(long, default_value = "main")]
+    pub(crate) revision: String,
+
+    /// Use flash attention
+    #[arg(long)]
+    pub(crate) use_flash_attn: bool,
+
+    /// Penalty to be applied for repeating tokens, 1. means no penalty
+    #[arg(long, default_value_t = 1.1)]
+    pub(crate) repeat_penalty: f32,
+
+    /// The context size to consider for the repeat penalty
+    #[arg(long, default_value_t = 64)]
+    pub(crate) repeat_last_n: usize,
+
+    /// Enable tracing
+    #[arg(long)]
+    pub(crate) tracing: bool,
+}
+
+pub fn run_cli() -> anyhow::Result<()> {
+    let args = Args::parse();
+    let cfg = GemmaInferenceConfig {
+        tracing: args.tracing,
+        prompt: args.prompt,
+        model: args.model,
+        cpu: args.cpu,
+        dtype: args.dtype,
+        model_id: args.model_id,
+        revision: args.revision,
+        use_flash_attn: args.use_flash_attn,
+        seed: args.seed,
+        temperature: args.temperature.unwrap_or(0.8),
+        top_p: args.top_p,
+        repeat_penalty: args.repeat_penalty,
+        repeat_last_n: args.repeat_last_n,
+        max_tokens: args.max_tokens,
+    };
+    let rx = run_gemma_api(cfg)?;
+    for msg in rx {
+        match msg {
+            Ok(tok) => {
+                print!("{tok}");
+                let _ = std::io::stdout().flush(); // <- force it out now
+            }
+            Err(e) => {
+                eprintln!("generation error: {e}");
+                break;
+            }
+        }
+    }
+    Ok(())
+}

crates/gemma-runner/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod gemma_api;
+
+pub use gemma_api::{run_gemma_api, GemmaInferenceConfig, WhichModel};

crates/gemma-runner/src/main.rs

@@ -0,0 +1,17 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+mod gemma_api;
+mod gemma_cli;
+
+use anyhow::Error;
+use clap::{Parser, ValueEnum};
+
+use crate::gemma_cli::run_cli;
+use std::io::Write;
+
+/// just a placeholder, not used for anything
+fn main() -> std::result::Result<(), Error> {
+    run_cli()
+}

crates/helm-chart-tool/Cargo.toml

@@ -3,8 +3,6 @@ name = "helm-chart-tool"
 version = "0.1.0"
 edition = "2021"
 
-[workspace]
-
 [[bin]]
 name = "helm-chart-tool"
 path = "src/main.rs"

crates/helm-chart-tool/src/main.rs

@@ -84,7 +84,10 @@ fn main() -> Result<()> {
     let services = discover_services(workspace_path)?;
     println!("Found {} services:", services.len());
     for service in &services {
-        println!("  - {}: {} (port {})", service.name, service.image, service.port);
+        println!(
+            "  - {}: {} (port {})",
+            service.name, service.image, service.port
+        );
     }
 
     generate_helm_chart(output_path, chart_name, &services)?;
@@ -115,17 +118,20 @@ fn discover_services(workspace_path: &str) -> Result<Vec<ServiceInfo>> {
 fn parse_cargo_toml(path: &Path) -> Result<ServiceInfo> {
     let content = fs::read_to_string(path)
         .with_context(|| format!("Failed to read Cargo.toml at {:?}", path))?;
-    
+
     let cargo_toml: CargoToml = toml::from_str(&content)
         .with_context(|| format!("Failed to parse Cargo.toml at {:?}", path))?;
 
-    let package = cargo_toml.package
+    let package = cargo_toml
+        .package
         .ok_or_else(|| anyhow::anyhow!("No package section found in {:?}", path))?;
 
-    let metadata = package.metadata
+    let metadata = package
+        .metadata
         .ok_or_else(|| anyhow::anyhow!("No metadata section found in {:?}", path))?;
 
-    let kube_metadata = metadata.kube
+    let kube_metadata = metadata
+        .kube
         .ok_or_else(|| anyhow::anyhow!("No kube metadata found in {:?}", path))?;
 
     Ok(ServiceInfo {
@@ -136,7 +142,11 @@ fn parse_cargo_toml(path: &Path) -> Result<ServiceInfo> {
     })
 }
 
-fn generate_helm_chart(output_path: &str, chart_name: &str, services: &[ServiceInfo]) -> Result<()> {
+fn generate_helm_chart(
+    output_path: &str,
+    chart_name: &str,
+    services: &[ServiceInfo],
+) -> Result<()> {
     let chart_dir = Path::new(output_path);
     let templates_dir = chart_dir.join("templates");
 
@@ -512,4 +522,4 @@ fn generate_helmignore(chart_dir: &Path) -> Result<()> {
 
     fs::write(chart_dir.join(".helmignore"), helmignore_content)?;
     Ok(())
-}
+}

crates/inference-engine/Cargo.toml

@@ -3,32 +3,13 @@ name = "inference-engine"
 version = "0.1.0"
 edition = "2021"
 
-[[bin]]
-name="cli"
-path = "src/cli_main.rs"
-
-
 [dependencies]
-accelerate-src = { version = "0.3.2", optional = true }
-candle-datasets = { version = "=0.9.1", optional = true }
-candle-nn = { version = "=0.9.1" }
-candle-transformers = { version = "=0.9.1" }
+candle-core = { git = "https://github.com/huggingface/candle.git" }
+candle-nn = { git = "https://github.com/huggingface/candle.git" }
+candle-transformers = { git = "https://github.com/huggingface/candle.git" }
 candle-flash-attn = { version = "=0.9.1", optional = true }
 candle-onnx = { version = "=0.9.1", optional = true }
 
-csv = "1.3.0"
-cudarc = { version = "0.16.3", features = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc", "f16", "cuda-version-from-build-system", "dynamic-linking"], default-features=false, optional = true }
-half = { version = "2.5.0", features = ["num-traits", "use-intrinsics", "rand_distr"], optional = true }
-hf-hub = { version = "0.4.1", features = ["tokio"] }
-image = { version = "0.25.2", default-features = false, features = ["jpeg", "png"] }
-intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"], optional = true }
-num-traits = { version = "0.2.15" }
-palette = { version = "0.7.6", optional = true }
-enterpolation = { version = "0.2.1", optional = true}
-pyo3 = { version = "0.22.0", features = ["auto-initialize", "abi3-py311"], optional = true }
-rayon = "1.7.0"
-rubato = { version = "0.15.0", optional = true }
-safetensors = "0.4.1"
 serde = { version = "1.0.171", features = ["derive"] }
 serde_json = "1.0.99"
 symphonia = { version = "0.5.3", features = ["all"], optional = true }
@@ -50,20 +31,14 @@ utoipa = { version = "4.2.0", features = ["axum_extras"] }
 uuid = { version = "1.7.0", features = ["v4"] }
 reborrow = "0.5.5"
 futures-util = "0.3.31"
+gemma-runner = { path = "../gemma-runner" }
+llama-runner = { path = "../llama-runner" }
 
-# --- Add this section for conditional compilation ---
 [target.'cfg(target_os = "macos")'.dependencies]
-# Use CPU backend for macOS to avoid Metal rotary-emb implementation issues
-candle-core = { version = "=0.9.1", features = ["metal"], optional = false }
+candle-core = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+candle-nn = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+candle-transformers = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
 
-[target.'cfg(not(target_os = "macos"))'.dependencies]
-# For Linux or other non-macOS systems, you likely want the CPU backend or CUDA
-# If you're building on Linux with a CUDA-enabled GPU:
-candle-core = { version = "=0.9.1", features = ["cuda"], default-features = false } # Or just "cuda" if not using default features
-
-# If you're building on Linux with only CPU:
-# candle-core = { version = "=0.9.1", default-features = false } # CPU is often the default, but good to be explicit
-# --- End of conditional compilation section ---
 
 [dev-dependencies]
 anyhow = { version = "1", features = ["backtrace"] }
@@ -83,6 +58,9 @@ tokio = "1.43.0"
 anyhow = { version = "1", features = ["backtrace"] }
 bindgen_cuda = { version = "0.1.1", optional = true }
 
+[features]
+bin = []
+
 
 
 [package.metadata.compose]

crates/inference-engine/src/cli.rs

@@ -1,72 +0,0 @@
-use clap::Parser;
-use crate::model::Which;
-
-#[derive(Parser, Debug)]
-#[command(author, version, about, long_about = None)]
-pub struct Args {
-    /// Run on CPU rather than on GPU.
-    #[arg(long)]
-    pub cpu: bool,
-
-    /// Enable tracing (generates a trace-timestamp.json file).
-    #[arg(long)]
-    pub tracing: bool,
-
-    /// Run in server mode with OpenAI compatible API
-    #[arg(long)]
-    pub server: bool,
-
-    /// Port to use for the server
-    #[arg(long, default_value_t = 3777)]
-    pub port: u16,
-
-    /// Prompt for text generation (not used in server mode)
-    #[arg(long)]
-    pub prompt: Option<String>,
-
-    /// The temperature used to generate samples.
-    #[arg(long)]
-    pub temperature: Option<f64>,
-
-    /// Nucleus sampling probability cutoff.
-    #[arg(long)]
-    pub top_p: Option<f64>,
-
-    /// The seed to use when generating random samples.
-    #[arg(long, default_value_t = 299792458)]
-    pub seed: u64,
-
-    /// The length of the sample to generate (in tokens).
-    #[arg(long, short = 'n', default_value_t = 10000)]
-    pub sample_len: usize,
-
-    #[arg(long)]
-    pub model_id: Option<String>,
-
-    #[arg(long, default_value = "main")]
-    pub revision: String,
-
-    #[arg(long)]
-    pub tokenizer_file: Option<String>,
-
-    #[arg(long)]
-    pub config_file: Option<String>,
-
-    #[arg(long)]
-    pub weight_files: Option<String>,
-
-    /// Penalty to be applied for repeating tokens, 1. means no penalty.
-    #[arg(long, default_value_t = 1.1)]
-    pub repeat_penalty: f32,
-
-    /// The context size to consider for the repeat penalty.
-    #[arg(long, default_value_t = 64)]
-    pub repeat_last_n: usize,
-
-    /// The model to use.
-    #[arg(long, default_value = "3-1b-it")]
-    pub which: Which,
-
-    #[arg(long)]
-    pub use_flash_attn: bool,
-}

crates/inference-engine/src/cli_main.rs

@@ -1,912 +0,0 @@
-mod token_output_stream;
-mod utilities_lib;
-
-#[cfg(feature = "intel-mkl-src")]
-extern crate intel_mkl_src;
-
-#[cfg(feature = "accelerate-src")]
-extern crate accelerate_src;
-
-#[cfg(feature = "metal")]
-extern crate metal_src;
-
-use anyhow::{Error as E, Result};
-use axum::{
-    extract::State,
-    http::StatusCode,
-    response::IntoResponse,
-    routing::{get, post},
-    Json, Router,
-};
-use clap::Parser;
-use either::Either;
-use serde::{Deserialize, Serialize};
-use std::{collections::HashMap, net::SocketAddr, sync::Arc};
-use tokio::sync::Mutex;
-use tower_http::cors::{Any, CorsLayer};
-use utoipa::ToSchema;
-
-use candle_transformers::models::gemma::{Config as Config1, Model as Model1};
-use candle_transformers::models::gemma2::{Config as Config2, Model as Model2};
-use candle_transformers::models::gemma3::{Config as Config3, Model as Model3};
-
-// OpenAI API compatible structs
-
-/// Inner content structure for messages that can be either a string or key-value pairs
-#[derive(Debug, Clone, Deserialize, Serialize)]
-pub struct MessageInnerContent(
-    #[serde(with = "either::serde_untagged")] pub Either<String, HashMap<String, String>>,
-);
-
-impl ToSchema<'_> for MessageInnerContent {
-    fn schema() -> (&'static str, utoipa::openapi::RefOr<utoipa::openapi::Schema>) {
-        (
-            "MessageInnerContent",
-            utoipa::openapi::RefOr::T(message_inner_content_schema()),
-        )
-    }
-}
-
-/// Function for MessageInnerContent Schema generation to handle `Either`
-fn message_inner_content_schema() -> utoipa::openapi::Schema {
-    use utoipa::openapi::{ArrayBuilder, ObjectBuilder, OneOfBuilder, RefOr, Schema, SchemaType};
-
-    Schema::OneOf(
-        OneOfBuilder::new()
-            // Either::Left - simple string
-            .item(Schema::Object(
-                ObjectBuilder::new().schema_type(SchemaType::String).build(),
-            ))
-            // Either::Right - object with string values
-            .item(Schema::Object(
-                ObjectBuilder::new()
-                    .schema_type(SchemaType::Object)
-                    .additional_properties(Some(RefOr::T(Schema::Object(
-                        ObjectBuilder::new().schema_type(SchemaType::String).build(),
-                    ))))
-                    .build(),
-            ))
-            .build(),
-    )
-}
-
-/// Message content that can be either simple text or complex structured content
-#[derive(Debug, Clone, Deserialize, Serialize)]
-pub struct MessageContent(
-    #[serde(with = "either::serde_untagged")]
-    Either<String, Vec<HashMap<String, MessageInnerContent>>>,
-);
-
-impl ToSchema<'_> for MessageContent {
-    fn schema() -> (&'static str, utoipa::openapi::RefOr<utoipa::openapi::Schema>) {
-        ("MessageContent", utoipa::openapi::RefOr::T(message_content_schema()))
-    }
-}
-
-/// Function for MessageContent Schema generation to handle `Either`
-fn message_content_schema() -> utoipa::openapi::Schema {
-    use utoipa::openapi::{ArrayBuilder, ObjectBuilder, OneOfBuilder, RefOr, Schema, SchemaType};
-
-    Schema::OneOf(
-        OneOfBuilder::new()
-            .item(Schema::Object(
-                ObjectBuilder::new().schema_type(SchemaType::String).build(),
-            ))
-            .item(Schema::Array(
-                ArrayBuilder::new()
-                    .items(RefOr::T(Schema::Object(
-                        ObjectBuilder::new()
-                            .schema_type(SchemaType::Object)
-                            .additional_properties(Some(RefOr::Ref(
-                                utoipa::openapi::Ref::from_schema_name("MessageInnerContent"),
-                            )))
-                            .build(),
-                    )))
-                    .build(),
-            ))
-            .build(),
-    )
-}
-
-/// Represents a single message in a conversation
-#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)]
-pub struct Message {
-    /// The message content
-    pub content: Option<MessageContent>,
-    /// The role of the message sender ("user", "assistant", "system", "tool", etc.)
-    pub role: String,
-    pub name: Option<String>,
-}
-
-/// Stop token configuration for generation
-#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)]
-#[serde(untagged)]
-pub enum StopTokens {
-    ///  Multiple possible stop sequences
-    Multi(Vec<String>),
-    /// Single stop sequence
-    Single(String),
-}
-
-/// Default value helper
-fn default_false() -> bool {
-    false
-}
-
-/// Default value helper
-fn default_1usize() -> usize {
-    1
-}
-
-/// Default value helper
-fn default_model() -> String {
-    "default".to_string()
-}
-
-/// Chat completion request following OpenAI's specification
-#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)]
-pub struct ChatCompletionRequest {
-    #[schema(example = json!([{"role": "user", "content": "Why did the crab cross the road?"}]))]
-    pub messages: Vec<Message>,
-    #[schema(example = "gemma-3-1b-it")]
-    #[serde(default = "default_model")]
-    pub model: String,
-    #[serde(default = "default_false")]
-    #[schema(example = false)]
-    pub logprobs: bool,
-    #[schema(example = 256)]
-    pub max_tokens: Option<usize>,
-    #[serde(rename = "n")]
-    #[serde(default = "default_1usize")]
-    #[schema(example = 1)]
-    pub n_choices: usize,
-    #[schema(example = 0.7)]
-    pub temperature: Option<f64>,
-    #[schema(example = 0.9)]
-    pub top_p: Option<f64>,
-    #[schema(example = false)]
-    pub stream: Option<bool>,
-}
-
-/// Chat completion choice
-#[derive(Debug, Serialize, ToSchema)]
-pub struct ChatCompletionChoice {
-    pub index: usize,
-    pub message: Message,
-    pub finish_reason: String,
-}
-
-/// Chat completion response
-#[derive(Debug, Serialize, ToSchema)]
-pub struct ChatCompletionResponse {
-    pub id: String,
-    pub object: String,
-    pub created: u64,
-    pub model: String,
-    pub choices: Vec<ChatCompletionChoice>,
-    pub usage: Usage,
-}
-
-/// Token usage information
-#[derive(Debug, Serialize, ToSchema)]
-pub struct Usage {
-    pub prompt_tokens: usize,
-    pub completion_tokens: usize,
-    pub total_tokens: usize,
-}
-
-// Application state shared between handlers
-#[derive(Clone)]
-struct AppState {
-    text_generation: Arc<Mutex<TextGeneration>>,
-    model_id: String,
-}
-
-// Chat completions endpoint handler
-async fn chat_completions(
-    State(state): State<AppState>,
-    Json(request): Json<ChatCompletionRequest>,
-) -> Result<Json<ChatCompletionResponse>, (StatusCode, Json<serde_json::Value>)> {
-    let mut prompt = String::new();
-
-    // Convert messages to a prompt string
-    for message in &request.messages {
-        let role = &message.role;
-        let content = match &message.content {
-            Some(content) => match &content.0 {
-                Either::Left(text) => text.clone(),
-                Either::Right(_) => "".to_string(), // Handle complex content if needed
-            },
-            None => "".to_string(),
-        };
-
-        // Format based on role
-        match role.as_str() {
-            "system" => prompt.push_str(&format!("System: {}\n", content)),
-            "user" => prompt.push_str(&format!("User: {}\n", content)),
-            "assistant" => prompt.push_str(&format!("Assistant: {}\n", content)),
-            _ => prompt.push_str(&format!("{}: {}\n", role, content)),
-        }
-    }
-
-    // Add the assistant prefix for the response
-    prompt.push_str("Assistant: ");
-
-    // Capture the output
-    let mut output = Vec::new();
-    {
-        let mut text_gen = state.text_generation.lock().await;
-
-        // Buffer to capture the output
-        let mut buffer = Vec::new();
-
-        // Run text generation
-        let max_tokens = request.max_tokens.unwrap_or(1000);
-        let result = text_gen.run_with_output(&prompt, max_tokens, &mut buffer);
-
-        if let Err(e) = result {
-    return Err((
-        StatusCode::BAD_REQUEST,
-        Json(serde_json::json!({
-            "error": {
-                "message": "The OpenAI API is currently not supported due to compatibility issues with the tensor operations. Please use the CLI mode instead with: cargo run --bin inference-engine -- --prompt \"Your prompt here\"",
-                "type": "unsupported_api"
-            }
-        })),
-    ));
-}
-
-        // Convert buffer to string
-        if let Ok(text) = String::from_utf8(buffer) {
-            output.push(text);
-        }
-    }
-
-    // Create response
-    let response = ChatCompletionResponse {
-        id: format!("chatcmpl-{}", uuid::Uuid::new_v4().to_string().replace("-", "")),
-        object: "chat.completion".to_string(),
-        created: std::time::SystemTime::now()
-            .duration_since(std::time::UNIX_EPOCH)
-            .unwrap_or_default()
-            .as_secs(),
-        model: request.model,
-        choices: vec![ChatCompletionChoice {
-            index: 0,
-            message: Message {
-                role: "assistant".to_string(),
-                content: Some(MessageContent(Either::Left(output.join("")))),
-                name: None,
-            },
-            finish_reason: "stop".to_string(),
-        }],
-        usage: Usage {
-            prompt_tokens: prompt.len() / 4, // Rough estimate
-            completion_tokens: output.join("").len() / 4, // Rough estimate
-            total_tokens: (prompt.len() + output.join("").len()) / 4, // Rough estimate
-        },
-    };
-
-    // Return the response as JSON
-    Ok(Json(response))
-}
-
-use candle_core::{DType, Device, MetalDevice, Tensor};
-use candle_nn::VarBuilder;
-use candle_transformers::generation::LogitsProcessor;
-use hf_hub::{Repo, RepoType, api::sync::Api};
-use serde_json::json;
-use tokenizers::Tokenizer;
-use crate::token_output_stream::TokenOutputStream;
-use crate::utilities_lib::device;
-
-// Create the router with the chat completions endpoint
-fn create_router(app_state: AppState) -> Router {
-    // CORS layer to allow requests from any origin
-    let cors = CorsLayer::new()
-        .allow_origin(Any)
-        .allow_methods(Any)
-        .allow_headers(Any);
-
-    Router::new()
-        // OpenAI compatible endpoints
-        .route("/v1/chat/completions", post(chat_completions))
-        // Add more endpoints as needed
-        .layer(cors)
-        .with_state(app_state)
-}
-
-#[derive(Clone, Debug, Copy, PartialEq, Eq, clap::ValueEnum)]
-enum Which {
-    #[value(name = "2b")]
-    Base2B,
-    #[value(name = "7b")]
-    Base7B,
-    #[value(name = "2b-it")]
-    Instruct2B,
-    #[value(name = "7b-it")]
-    Instruct7B,
-    #[value(name = "1.1-2b-it")]
-    InstructV1_1_2B,
-    #[value(name = "1.1-7b-it")]
-    InstructV1_1_7B,
-    #[value(name = "code-2b")]
-    CodeBase2B,
-    #[value(name = "code-7b")]
-    CodeBase7B,
-    #[value(name = "code-2b-it")]
-    CodeInstruct2B,
-    #[value(name = "code-7b-it")]
-    CodeInstruct7B,
-    #[value(name = "2-2b")]
-    BaseV2_2B,
-    #[value(name = "2-2b-it")]
-    InstructV2_2B,
-    #[value(name = "2-9b")]
-    BaseV2_9B,
-    #[value(name = "2-9b-it")]
-    InstructV2_9B,
-    #[value(name = "3-1b")]
-    BaseV3_1B,
-    #[value(name = "3-1b-it")]
-    InstructV3_1B,
-}
-
-enum Model {
-    V1(Model1),
-    V2(Model2),
-    V3(Model3),
-}
-
-impl Model {
-    fn forward(&mut self, input_ids: &candle_core::Tensor, pos: usize) -> candle_core::Result<candle_core::Tensor> {
-        match self {
-            Self::V1(m) => m.forward(input_ids, pos),
-            Self::V2(m) => m.forward(input_ids, pos),
-            Self::V3(m) => m.forward(input_ids, pos),
-        }
-    }
-}
-
-
-
-struct TextGeneration {
-    model: Model,
-    device: Device,
-    tokenizer: TokenOutputStream,
-    logits_processor: LogitsProcessor,
-    repeat_penalty: f32,
-    repeat_last_n: usize,
-}
-
-impl TextGeneration {
-    #[allow(clippy::too_many_arguments)]
-    fn new(
-        model: Model,
-        tokenizer: Tokenizer,
-        seed: u64,
-        temp: Option<f64>,
-        top_p: Option<f64>,
-        repeat_penalty: f32,
-        repeat_last_n: usize,
-        device: &Device,
-    ) -> Self {
-        let logits_processor = LogitsProcessor::new(seed, temp, top_p);
-        Self {
-            model,
-            tokenizer: TokenOutputStream::new(tokenizer),
-            logits_processor,
-            repeat_penalty,
-            repeat_last_n,
-            device: device.clone(),
-        }
-    }
-
-    // Run text generation and print to stdout
-    fn run(&mut self, prompt: &str, sample_len: usize) -> Result<()> {
-        use std::io::Write;
-        self.tokenizer.clear();
-        let mut tokens = self
-            .tokenizer
-            .tokenizer()
-            .encode(prompt, true)
-            .map_err(E::msg)?
-            .get_ids()
-            .to_vec();
-        for &t in tokens.iter() {
-            if let Some(t) = self.tokenizer.next_token(t)? {
-                print!("{t}")
-            }
-        }
-        std::io::stdout().flush()?;
-
-        let mut generated_tokens = 0usize;
-        let eos_token = match self.tokenizer.get_token("<eos>") {
-            Some(token) => token,
-            None => anyhow::bail!("cannot find the <eos> token"),
-        };
-
-        let eot_token = match self.tokenizer.get_token("<end_of_turn>") {
-            Some(token) => token,
-            None => {
-                println!(
-                    "Warning: <end_of_turn> token not found in tokenizer, using <eos> as a backup"
-                );
-                eos_token
-            }
-        };
-
-        let start_gen = std::time::Instant::now();
-        for index in 0..sample_len {
-            let context_size = if index > 0 { 1 } else { tokens.len() };
-            let start_pos = tokens.len().saturating_sub(context_size);
-            let ctxt = &tokens[start_pos..];
-            let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
-            let logits = self.model.forward(&input, start_pos)?;
-            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
-            let logits = if self.repeat_penalty == 1. {
-                logits
-            } else {
-                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
-
-                // Manual implementation of repeat penalty to avoid type conflicts
-                let mut logits_vec = logits.to_vec1::<f32>()?;
-
-                for &token_id in &tokens[start_at..] {
-                    let token_id = token_id as usize;
-                    if token_id < logits_vec.len() {
-                        let score = logits_vec[token_id];
-                        let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                        logits_vec[token_id] = sign * score / self.repeat_penalty;
-                    }
-                }
-
-                // Create a new tensor with the modified logits
-                let device = logits.device().clone();
-                let shape = logits.shape().clone();
-                let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                new_logits.reshape(shape)?
-            };
-
-            let next_token = self.logits_processor.sample(&logits)?;
-            tokens.push(next_token);
-            generated_tokens += 1;
-            if next_token == eos_token || next_token == eot_token {
-                break;
-            }
-            if let Some(t) = self.tokenizer.next_token(next_token)? {
-                print!("{t}");
-                std::io::stdout().flush()?;
-            }
-        }
-        let dt = start_gen.elapsed();
-        if let Some(rest) = self.tokenizer.decode_rest().map_err(E::msg)? {
-            print!("{rest}");
-        }
-        std::io::stdout().flush()?;
-        println!(
-            "\n{generated_tokens} tokens generated ({:.2} token/s)",
-            generated_tokens as f64 / dt.as_secs_f64(),
-        );
-        Ok(())
-    }
-
-    // Run text generation and write to a buffer
-    fn run_with_output(&mut self, prompt: &str, sample_len: usize, output: &mut Vec<u8>) -> Result<()> {
-        use std::io::Write;
-        self.tokenizer.clear();
-        let mut tokens = self
-            .tokenizer
-            .tokenizer()
-            .encode(prompt, true)
-            .map_err(E::msg)?
-            .get_ids()
-            .to_vec();
-
-        // Write prompt tokens to output
-        for &t in tokens.iter() {
-            if let Some(t) = self.tokenizer.next_token(t)? {
-                write!(output, "{}", t)?;
-            }
-        }
-
-        let mut generated_tokens = 0usize;
-        let eos_token = match self.tokenizer.get_token("<eos>") {
-            Some(token) => token,
-            None => anyhow::bail!("cannot find the <eos> token"),
-        };
-
-        let eot_token = match self.tokenizer.get_token("<end_of_turn>") {
-            Some(token) => token,
-            None => {
-                write!(output, "Warning: <end_of_turn> token not found in tokenizer, using <eos> as a backup")?;
-                eos_token
-            }
-        };
-
-        // Determine if we're using a Model3 (gemma-3) variant
-        let is_model3 = match &self.model {
-            Model::V3(_) => true,
-            _ => false,
-        };
-
-        // For Model3, we need to use a different approach
-        if is_model3 {
-            // For gemma-3 models, we'll generate one token at a time with the full context
-            let start_gen = std::time::Instant::now();
-
-            // Initial generation with the full prompt
-            let input = Tensor::new(tokens.as_slice(), &self.device)?.unsqueeze(0)?;
-            let mut logits = self.model.forward(&input, 0)?;
-            logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
-
-            for _ in 0..sample_len {
-                // Apply repeat penalty if needed
-                let current_logits = if self.repeat_penalty == 1. {
-                    logits.clone()
-                } else {
-                    let start_at = tokens.len().saturating_sub(self.repeat_last_n);
-
-                    // Manual implementation of repeat penalty to avoid type conflicts
-                    let mut logits_vec = logits.to_vec1::<f32>()?;
-
-                    for &token_id in &tokens[start_at..] {
-                        let token_id = token_id as usize;
-                        if token_id < logits_vec.len() {
-                            let score = logits_vec[token_id];
-                            let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                            logits_vec[token_id] = sign * score / self.repeat_penalty;
-                        }
-                    }
-
-                    // Create a new tensor with the modified logits
-                    let device = logits.device().clone();
-                    let shape = logits.shape().clone();
-                    let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                    new_logits.reshape(shape)?
-                };
-
-                let next_token = self.logits_processor.sample(&current_logits)?;
-                tokens.push(next_token);
-                generated_tokens += 1;
-
-                if next_token == eos_token || next_token == eot_token {
-                    break;
-                }
-
-                if let Some(t) = self.tokenizer.next_token(next_token)? {
-                    write!(output, "{}", t)?;
-                }
-
-                // For the next iteration, just use the new token
-                let new_input = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?;
-                logits = self.model.forward(&new_input, tokens.len() - 1)?;
-                logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
-            }
-
-            return Ok(());
-        }
-
-        // Standard approach for other models
-        let start_gen = std::time::Instant::now();
-        for index in 0..sample_len {
-            let context_size = if index > 0 { 1 } else { tokens.len() };
-            let start_pos = tokens.len().saturating_sub(context_size);
-            let ctxt = &tokens[start_pos..];
-            let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
-            let logits = self.model.forward(&input, start_pos)?;
-            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
-            let logits = if self.repeat_penalty == 1. {
-                logits
-            } else {
-                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
- 
-                // Manual implementation of repeat penalty to avoid type conflicts
-                let mut logits_vec = logits.to_vec1::<f32>()?;
-
-                for &token_id in &tokens[start_at..] {
-                    let token_id = token_id as usize;
-                    if token_id < logits_vec.len() {
-                        let score = logits_vec[token_id];
-                        let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                        logits_vec[token_id] = sign * score / self.repeat_penalty;
-                    }
-                }
-
-                // Create a new tensor with the modified logits
-                let device = logits.device().clone();
-                let shape = logits.shape().clone();
-                let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                new_logits.reshape(shape)?
-            };
-
-            let next_token = self.logits_processor.sample(&logits)?;
-            tokens.push(next_token);
-            generated_tokens += 1;
-            if next_token == eos_token || next_token == eot_token {
-                break;
-            }
-            if let Some(t) = self.tokenizer.next_token(next_token)? {
-                write!(output, "{}", t)?;
-            }
-        }
-
-        // Write any remaining tokens
-        if let Some(rest) = self.tokenizer.decode_rest().map_err(E::msg)? {
-            write!(output, "{}", rest)?;
-        }
-
-        Ok(())
-    }
-}
-
-#[derive(Parser, Debug)]
-#[command(author, version, about, long_about = None)]
-struct Args {
-    /// Run on CPU rather than on GPU.
-    #[arg(long)]
-    cpu: bool,
-
-    /// Enable tracing (generates a trace-timestamp.json file).
-    #[arg(long)]
-    tracing: bool,
-
-    /// Run in server mode with OpenAI compatible API
-    #[arg(long)]
-    server: bool,
-
-    /// Port to use for the server
-    #[arg(long, default_value_t = 3777)]
-    port: u16,
-
-    /// Prompt for text generation (not used in server mode)
-    #[arg(long)]
-    prompt: Option<String>,
-
-    /// The temperature used to generate samples.
-    #[arg(long)]
-    temperature: Option<f64>,
-
-    /// Nucleus sampling probability cutoff.
-    #[arg(long)]
-    top_p: Option<f64>,
-
-    /// The seed to use when generating random samples.
-    #[arg(long, default_value_t = 299792458)]
-    seed: u64,
-
-    /// The length of the sample to generate (in tokens).
-    #[arg(long, short = 'n', default_value_t = 10000)]
-    sample_len: usize,
-
-    #[arg(long)]
-    model_id: Option<String>,
-
-    #[arg(long, default_value = "main")]
-    revision: String,
-
-    #[arg(long)]
-    tokenizer_file: Option<String>,
-
-    #[arg(long)]
-    config_file: Option<String>,
-
-    #[arg(long)]
-    weight_files: Option<String>,
-
-    /// Penalty to be applied for repeating tokens, 1. means no penalty.
-    #[arg(long, default_value_t = 1.1)]
-    repeat_penalty: f32,
-
-    /// The context size to consider for the repeat penalty.
-    #[arg(long, default_value_t = 64)]
-    repeat_last_n: usize,
-
-    /// The model to use.
-    #[arg(long, default_value = "3-1b-it")]
-    which: Which,
-
-    #[arg(long)]
-    use_flash_attn: bool,
-}
-
-fn main() -> Result<()> {
-    use tracing_chrome::ChromeLayerBuilder;
-    use tracing_subscriber::prelude::*;
-
-    let args = Args::parse();
-    let _guard = if args.tracing {
-        let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
-        tracing_subscriber::registry().with(chrome_layer).init();
-        Some(guard)
-    } else {
-        None
-    };
-    println!(
-        "avx: {}, neon: {}, simd128: {}, f16c: {}",
-        candle_core::utils::with_avx(),
-        candle_core::utils::with_neon(),
-        candle_core::utils::with_simd128(),
-        candle_core::utils::with_f16c()
-    );
-    println!(
-        "temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}",
-        args.temperature.unwrap_or(0.),
-        args.repeat_penalty,
-        args.repeat_last_n
-    );
-
-    let start = std::time::Instant::now();
-    let api = Api::new()?;
-    let model_id = match &args.model_id {
-        Some(model_id) => model_id.to_string(),
-        None => match args.which {
-            Which::InstructV1_1_2B => "google/gemma-1.1-2b-it".to_string(),
-            Which::InstructV1_1_7B => "google/gemma-1.1-7b-it".to_string(),
-            Which::Base2B => "google/gemma-2b".to_string(),
-            Which::Base7B => "google/gemma-7b".to_string(),
-            Which::Instruct2B => "google/gemma-2b-it".to_string(),
-            Which::Instruct7B => "google/gemma-7b-it".to_string(),
-            Which::CodeBase2B => "google/codegemma-2b".to_string(),
-            Which::CodeBase7B => "google/codegemma-7b".to_string(),
-            Which::CodeInstruct2B => "google/codegemma-2b-it".to_string(),
-            Which::CodeInstruct7B => "google/codegemma-7b-it".to_string(),
-            Which::BaseV2_2B => "google/gemma-2-2b".to_string(),
-            Which::InstructV2_2B => "google/gemma-2-2b-it".to_string(),
-            Which::BaseV2_9B => "google/gemma-2-9b".to_string(),
-            Which::InstructV2_9B => "google/gemma-2-9b-it".to_string(),
-            Which::BaseV3_1B => "google/gemma-3-1b-pt".to_string(),
-            Which::InstructV3_1B => "google/gemma-3-1b-it".to_string(),
-        },
-    };
-    let repo = api.repo(Repo::with_revision(
-        model_id.clone(),
-        RepoType::Model,
-        args.revision,
-    ));
-    let tokenizer_filename = match args.tokenizer_file {
-        Some(file) => std::path::PathBuf::from(file),
-        None => repo.get("tokenizer.json")?,
-    };
-    let config_filename = match args.config_file {
-        Some(file) => std::path::PathBuf::from(file),
-        None => repo.get("config.json")?,
-    };
-    let filenames = match args.weight_files {
-        Some(files) => files
-            .split(',')
-            .map(std::path::PathBuf::from)
-            .collect::<Vec<_>>(),
-        None => match args.which {
-            Which::BaseV3_1B | Which::InstructV3_1B => vec![repo.get("model.safetensors")?],
-            _ => utilities_lib::hub_load_safetensors(&repo, "model.safetensors.index.json")?,
-        },
-    };
-    println!("retrieved the files in {:?}", start.elapsed());
-    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
-
-    let start = std::time::Instant::now();
-    let initial_device = utilities_lib::device(args.cpu)?;
-    
-    // Check if we're using a V3 model (Gemma 3) and if we're on Metal (macOS)
-    let is_v3_model = matches!(args.which, Which::BaseV3_1B | Which::InstructV3_1B);
-    let is_metal = !initial_device.is_cpu() && candle_core::utils::metal_is_available() && !args.cpu;
-    
-    // Use CPU for V3 models on Metal due to missing implementations
-    let device = if is_v3_model && is_metal {
-        println!("Note: Using CPU for Gemma 3 model due to missing Metal implementations for required operations (e.g., rotary-emb).");
-        Device::Cpu
-    } else {
-        initial_device
-    };
-    
-    let dtype = if device.is_cuda() {
-        DType::BF16
-    } else {
-        DType::F32
-    };
-    
-    // Use the selected device and dtype
-    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, dtype, &device)? };
-    let model = match args.which {
-        Which::Base2B
-        | Which::Base7B
-        | Which::Instruct2B
-        | Which::Instruct7B
-        | Which::InstructV1_1_2B
-        | Which::InstructV1_1_7B
-        | Which::CodeBase2B
-        | Which::CodeBase7B
-        | Which::CodeInstruct2B
-        | Which::CodeInstruct7B => {
-            let config: Config1 = serde_json::from_reader(std::fs::File::open(config_filename)?)?;
-            let model = Model1::new(args.use_flash_attn, &config, vb)?;
-            Model::V1(model)
-        }
-        Which::BaseV2_2B | Which::InstructV2_2B | Which::BaseV2_9B | Which::InstructV2_9B => {
-            let config: Config2 = serde_json::from_reader(std::fs::File::open(config_filename)?)?;
-            let model = Model2::new(args.use_flash_attn, &config, vb)?;
-            Model::V2(model)
-        }
-        Which::BaseV3_1B | Which::InstructV3_1B => {
-            let config: Config3 = serde_json::from_reader(std::fs::File::open(config_filename)?)?;
-            let model = Model3::new(args.use_flash_attn, &config, vb)?;
-            Model::V3(model)
-        }
-    };
-
-    println!("loaded the model in {:?}", start.elapsed());
-
-    let pipeline = TextGeneration::new(
-        model,
-        tokenizer,
-        args.seed,
-        args.temperature,
-        args.top_p,
-        args.repeat_penalty,
-        args.repeat_last_n,
-        &device,
-    );
-
-    if args.server {
-        // Start the server
-        println!("Starting server on port {}", args.port);
-
-        // Create app state
-        let app_state = AppState {
-            text_generation: Arc::new(Mutex::new(pipeline)),
-            model_id,
-        };
-
-        // Create router
-        let app = create_router(app_state);
-
-        // Run the server
-        let addr = SocketAddr::from(([0, 0, 0, 0], args.port));
-
-        // Use tokio to run the server
-        tokio::runtime::Builder::new_multi_thread()
-            .enable_all()
-            .build()?
-            .block_on(async {
-                axum::serve(tokio::net::TcpListener::bind(&addr).await?, app)
-                    .await
-                    .map_err(|e| anyhow::anyhow!("Server error: {}", e))
-            })?;
-
-        Ok(())
-    } else {
-        // Run in CLI mode
-        if let Some(prompt_text) = &args.prompt {
-            let prompt = match args.which {
-                Which::Base2B
-                | Which::Base7B
-                | Which::Instruct2B
-                | Which::Instruct7B
-                | Which::InstructV1_1_2B
-                | Which::InstructV1_1_7B
-                | Which::CodeBase2B
-                | Which::CodeBase7B
-                | Which::CodeInstruct2B
-                | Which::CodeInstruct7B
-                | Which::BaseV2_2B
-                | Which::InstructV2_2B
-                | Which::BaseV2_9B
-                | Which::InstructV2_9B
-                | Which::BaseV3_1B => prompt_text.clone(),
-                Which::InstructV3_1B => {
-                    format!(
-                        "<start_of_turn> user\n{}<end_of_turn>\n<start_of_turn> model\n",
-                        prompt_text
-                    )
-                }
-            };
-
-            let mut pipeline = pipeline;
-            pipeline.run(&prompt, args.sample_len)?;
-            Ok(())
-        } else {
-            anyhow::bail!("Prompt is required in CLI mode. Use --prompt to specify a prompt or --server to run in server mode.")
-        }
-    }
-}

crates/inference-engine/src/inference.rs

@@ -0,0 +1,33 @@
+use anyhow::Result;
+use candle_core::Tensor;
+
+/// ModelInference trait defines the common interface for model inference operations
+///
+/// This trait serves as an abstraction for different model implementations (Gemma and Llama)
+/// to provide a unified interface for the inference engine.
+pub trait ModelInference {
+    /// Perform model inference for the given input tensor starting at the specified position
+    ///
+    /// # Arguments
+    ///
+    /// * `input_ids` - The input tensor containing token IDs
+    /// * `pos` - The position to start generation from
+    ///
+    /// # Returns
+    ///
+    /// A tensor containing the logits for the next token prediction
+    fn forward(&mut self, input_ids: &Tensor, pos: usize) -> Result<Tensor>;
+
+    /// Reset the model's internal state, if applicable
+    ///
+    /// This method can be used to clear any cached state between inference requests
+    fn reset_state(&mut self) -> Result<()>;
+
+    /// Get the model type name
+    ///
+    /// Returns a string identifier for the model type (e.g., "Gemma", "Llama")
+    fn model_type(&self) -> &'static str;
+}
+
+/// Factory function type for creating model inference implementations
+pub type ModelInferenceFactory = fn() -> Result<Box<dyn ModelInference>>;

crates/inference-engine/src/lib.rs

@@ -1,17 +1,14 @@
 // Expose modules for testing and library usage
-pub mod token_output_stream;
 pub mod model;
-pub mod text_generation;
-pub mod utilities_lib;
 pub mod openai_types;
-pub mod cli;
+// pub mod cli;
+pub mod inference;
 pub mod server;
 
 // Re-export key components for easier access
+pub use inference::ModelInference;
 pub use model::{Model, Which};
-pub use text_generation::TextGeneration;
-pub use token_output_stream::TokenOutputStream;
-pub use server::{AppState, create_router};
+pub use server::{create_router, AppState};
 
 use std::env;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

crates/inference-engine/src/model.rs

@@ -1,4 +1,5 @@
 // use candle_core::Tensor;
+use candle_transformers::models::csm::{LlamaConfig, LlamaModel};
 use candle_transformers::models::gemma::{Config as Config1, Model as Model1};
 use candle_transformers::models::gemma2::{Config as Config2, Model as Model2};
 use candle_transformers::models::gemma3::{Config as Config3, Model as Model3};
@@ -37,20 +38,30 @@ pub enum Which {
     BaseV3_1B,
     #[value(name = "3-1b-it")]
     InstructV3_1B,
+    #[value(name = "llama-3.2-1b-it")]
+    LlamaInstruct3_2_1B,
+    #[value(name = "llama-3.2-3b-it")]
+    LlamaInstruct3_2_3B,
 }
 
 pub enum Model {
     V1(Model1),
     V2(Model2),
     V3(Model3),
+    Llama(LlamaModel),
 }
 
 impl Model {
-    pub fn forward(&mut self, input_ids: &candle_core::Tensor, pos: usize) -> candle_core::Result<candle_core::Tensor> {
+    pub fn forward(
+        &mut self,
+        input_ids: &candle_core::Tensor,
+        pos: usize,
+    ) -> candle_core::Result<candle_core::Tensor> {
         match self {
             Self::V1(m) => m.forward(input_ids, pos),
             Self::V2(m) => m.forward(input_ids, pos),
             Self::V3(m) => m.forward(input_ids, pos),
+            Self::Llama(m) => m.forward(input_ids, pos),
         }
     }
 }
@@ -74,12 +85,20 @@ impl Which {
             Self::InstructV2_9B => "google/gemma-2-9b-it".to_string(),
             Self::BaseV3_1B => "google/gemma-3-1b-pt".to_string(),
             Self::InstructV3_1B => "google/gemma-3-1b-it".to_string(),
+            Self::LlamaInstruct3_2_1B => "meta-llama/Llama-3.2-1B-Instruct".to_string(),
+            Self::LlamaInstruct3_2_3B => "meta-llama/Llama-3.2-3B-Instruct".to_string(),
         }
     }
 
     pub fn is_instruct_model(&self) -> bool {
         match self {
-            Self::Base2B | Self::Base7B | Self::CodeBase2B | Self::CodeBase7B | Self::BaseV2_2B | Self::BaseV2_9B | Self::BaseV3_1B => false,
+            Self::Base2B
+            | Self::Base7B
+            | Self::CodeBase2B
+            | Self::CodeBase7B
+            | Self::BaseV2_2B
+            | Self::BaseV2_9B
+            | Self::BaseV3_1B => false,
             _ => true,
         }
     }
@@ -87,4 +106,8 @@ impl Which {
     pub fn is_v3_model(&self) -> bool {
         matches!(self, Self::BaseV3_1B | Self::InstructV3_1B)
     }
-}
+
+    pub fn is_llama_model(&self) -> bool {
+        matches!(self, Self::LlamaInstruct3_2_1B | Self::LlamaInstruct3_2_3B)
+    }
+}

crates/inference-engine/src/openai_types.rs

@@ -1,5 +1,6 @@
 use either::Either;
 use serde::{Deserialize, Serialize};
+use serde_json::json;
 use std::collections::HashMap;
 use utoipa::ToSchema;
 
@@ -10,7 +11,10 @@ pub struct MessageInnerContent(
 );
 
 impl ToSchema<'_> for MessageInnerContent {
-    fn schema() -> (&'static str, utoipa::openapi::RefOr<utoipa::openapi::Schema>) {
+    fn schema() -> (
+        &'static str,
+        utoipa::openapi::RefOr<utoipa::openapi::Schema>,
+    ) {
         (
             "MessageInnerContent",
             utoipa::openapi::RefOr::T(message_inner_content_schema()),
@@ -45,12 +49,18 @@ fn message_inner_content_schema() -> utoipa::openapi::Schema {
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct MessageContent(
     #[serde(with = "either::serde_untagged")]
-    pub Either<String, Vec<HashMap<String, MessageInnerContent>>>,
+    pub  Either<String, Vec<HashMap<String, MessageInnerContent>>>,
 );
 
 impl ToSchema<'_> for MessageContent {
-    fn schema() -> (&'static str, utoipa::openapi::RefOr<utoipa::openapi::Schema>) {
-        ("MessageContent", utoipa::openapi::RefOr::T(message_content_schema()))
+    fn schema() -> (
+        &'static str,
+        utoipa::openapi::RefOr<utoipa::openapi::Schema>,
+    ) {
+        (
+            "MessageContent",
+            utoipa::openapi::RefOr::T(message_content_schema()),
+        )
     }
 }
 
@@ -213,4 +223,4 @@ pub struct ModelListResponse {
     pub object: String,
     /// Array of available models
     pub data: Vec<Model>,
-}
+}

crates/inference-engine/src/server.rs

@@ -5,299 +5,86 @@ use axum::{
     routing::{get, post},
     Json, Router,
 };
-use candle_core::DType;
-use candle_nn::VarBuilder;
 use futures_util::stream::{self, Stream};
-use tokio_stream::wrappers::UnboundedReceiverStream;
 use std::convert::Infallible;
-use std::{path::PathBuf, sync::Arc};
-use tokio::sync::{Mutex, mpsc};
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio_stream::wrappers::UnboundedReceiverStream;
 use tower_http::cors::{Any, CorsLayer};
 use uuid::Uuid;
 
-use crate::openai_types::{ChatCompletionChoice, ChatCompletionChunk, ChatCompletionChunkChoice, ChatCompletionRequest, ChatCompletionResponse, Delta, Message, MessageContent, Model, ModelListResponse, Usage};
-use crate::text_generation::TextGeneration;
-use crate::{utilities_lib, Model as GemmaModel, Which};
+use crate::openai_types::{
+    ChatCompletionChoice, ChatCompletionChunk, ChatCompletionChunkChoice, ChatCompletionRequest,
+    ChatCompletionResponse, Delta, Message, MessageContent, Model, ModelListResponse, Usage,
+};
+use crate::Which;
 use either::Either;
-use hf_hub::api::sync::{Api, ApiError};
-use hf_hub::{Repo, RepoType};
-use tokenizers::Tokenizer;
-
-use candle_transformers::models::gemma::{Config as Config1, Model as Model1};
-use candle_transformers::models::gemma2::{Config as Config2, Model as Model2};
-use candle_transformers::models::gemma3::{Config as Config3, Model as Model3};
+use gemma_runner::{run_gemma_api, GemmaInferenceConfig};
+use llama_runner::{run_llama_inference, LlamaInferenceConfig};
 use serde_json::Value;
 // -------------------------
 // Shared app state
 // -------------------------
 
+#[derive(Clone, Debug)]
+pub enum ModelType {
+    Gemma,
+    Llama,
+}
+
 #[derive(Clone)]
 pub struct AppState {
-    pub text_generation: Arc<Mutex<TextGeneration>>,
+    pub model_type: ModelType,
     pub model_id: String,
-    // Store build args to recreate TextGeneration when needed
-    pub build_args: PipelineArgs,
+    pub gemma_config: Option<GemmaInferenceConfig>,
+    pub llama_config: Option<LlamaInferenceConfig>,
 }
 
 impl Default for AppState {
     fn default() -> Self {
-        let args = PipelineArgs::default();
-        let text_generation = build_pipeline(args.clone());
+        let gemma_config = GemmaInferenceConfig {
+            model: gemma_runner::WhichModel::InstructV3_1B,
+            ..Default::default()
+        };
         Self {
-            text_generation: Arc::new(Mutex::new(text_generation)),
-            model_id: args.model_id.clone(),
-            build_args: args,
+            model_type: ModelType::Gemma,
+            model_id: "gemma-3-1b-it".to_string(),
+            gemma_config: Some(gemma_config),
+            llama_config: None,
         }
     }
 }
 
 // -------------------------
-// Pipeline configuration
+// Helper functions
 // -------------------------
 
-#[derive(Debug, Clone)]
-pub struct PipelineArgs {
-    pub model_id: String,
-    pub which: Which,
-    pub revision: Option<String>,
-    pub tokenizer_path: Option<PathBuf>,
-    pub config_path: Option<PathBuf>,
-    pub weight_paths: Vec<PathBuf>,
-    pub use_flash_attn: bool,
-    pub force_cpu: bool,
-    pub seed: u64,
-    pub temperature: Option<f64>,
-    pub top_p: Option<f64>,
-    pub repeat_penalty: f32,
-    pub repeat_last_n: usize,
-}
-
-impl Default for PipelineArgs {
-    fn default() -> Self {
-        Self {
-            model_id: Which::InstructV3_1B.to_model_id().to_string(),
-            which: Which::InstructV3_1B,
-            revision: None,
-            tokenizer_path: None,
-            config_path: None,
-            weight_paths: Vec::new(),
-            use_flash_attn: false,
-            force_cpu: false,
-            seed: 299792458, // Speed of light in vacuum (m/s)
-            temperature: Some(0.8), // Good balance between creativity and coherence
-            top_p: Some(0.9), // Keep diverse but reasonable options
-            repeat_penalty: 1.2, // Stronger penalty for repetition to prevent looping
-            repeat_last_n: 64, // Consider last 64 tokens for repetition
-        }
-    }
-}
-
 fn normalize_model_id(model_id: &str) -> String {
-    if model_id.contains('/') {
-        model_id.to_string()
-    } else {
-        format!("google/{}", model_id)
-    }
-}
-
-fn ensure_repo_exists(api: &Api, model_id: &str, revision: &str) -> anyhow::Result<()> {
-    let repo = api.repo(Repo::with_revision(
-        model_id.to_string(),
-        RepoType::Model,
-        revision.to_string(),
-    ));
-    match repo.get("config.json") {
-        Ok(_) => Ok(()),
-        Err(e) => match e {
-            ApiError::RequestError(resp) => {
-                let error_str = resp.to_string();
-                if error_str.contains("404") {
-                    anyhow::bail!(
-                        "Hugging Face model repo not found: '{model_id}' at revision '{revision}'."
-                    )
-                }
-                Err(anyhow::Error::new(ApiError::RequestError(resp)))
-            }
-            other => Err(anyhow::Error::new(other)),
-        },
-    }
-}
-
-// -------------------------
-// Pipeline builder
-// -------------------------
-
-pub fn build_pipeline(mut args: PipelineArgs) -> TextGeneration {
-    println!(
-        "avx: {}, neon: {}, simd128: {}, f16c: {}",
-        candle_core::utils::with_avx(),
-        candle_core::utils::with_neon(),
-        candle_core::utils::with_simd128(),
-        candle_core::utils::with_f16c()
-    );
-
-    let start = std::time::Instant::now();
-    let api = Api::new().unwrap();
-    let revision = args.revision.as_deref().unwrap_or("main");
-
-    if args.model_id.trim().is_empty() {
-        panic!("No model ID specified.");
-    }
-    args.model_id = normalize_model_id(&args.model_id);
-
-    match ensure_repo_exists(&api, &args.model_id, revision) {
-        Ok(_) => {}
-        Err(e) => panic!("{}", e),
-    };
-
-    let repo = api.repo(Repo::with_revision(
-        args.model_id.clone(),
-        RepoType::Model,
-        revision.to_string(),
-    ));
-
-    let tokenizer_path = args
-        .tokenizer_path
-        .unwrap_or_else(|| repo.get("tokenizer.json").unwrap());
-    let config_path = args
-        .config_path
-        .unwrap_or_else(|| repo.get("config.json").unwrap());
-
-    if !matches!(
-        args.which,
-        Which::Base2B
-            | Which::Base7B
-            | Which::Instruct2B
-            | Which::Instruct7B
-            | Which::InstructV1_1_2B
-            | Which::InstructV1_1_7B
-            | Which::CodeBase2B
-            | Which::CodeBase7B
-            | Which::CodeInstruct2B
-            | Which::CodeInstruct7B
-            | Which::BaseV2_2B
-            | Which::InstructV2_2B
-            | Which::BaseV2_9B
-            | Which::InstructV2_9B
-            | Which::BaseV3_1B
-            | Which::InstructV3_1B
-    ) {
-        if args.model_id.contains("gemma-2-2b-it") {
-            args.which = Which::InstructV2_2B;
-        } else if args.model_id.contains("gemma-3-1b-it") {
-            args.which = Which::InstructV3_1B;
-        } else if let Ok(file) = std::fs::File::open(config_path.clone()) {
-            if let Ok(cfg_val) = serde_json::from_reader::<_, serde_json::Value>(file) {
-                if let Some(model_type) = cfg_val.get("model_type").and_then(|v| v.as_str()) {
-                    if model_type.contains("gemma3") {
-                        args.which = Which::InstructV3_1B;
-                    } else if model_type.contains("gemma2") {
-                        args.which = Which::InstructV2_2B;
-                    } else {
-                        args.which = Which::Instruct2B;
-                    }
-                }
-            }
-        }
-    }
-
-    let weight_paths = if !args.weight_paths.is_empty() {
-        args.weight_paths
-    } else {
-        match repo.get("model.safetensors") {
-            Ok(single) => vec![single],
-            Err(_) => match utilities_lib::hub_load_safetensors(&repo, "model.safetensors.index.json") {
-                Ok(paths) => paths,
-                Err(e) => {
-                    panic!("Unable to locate model weights: {}", e);
-                }
-            },
-        }
-    };
-
-    println!("retrieved the files in {:?}", start.elapsed());
-
-    let tokenizer = Tokenizer::from_file(tokenizer_path).unwrap();
-
-    let initial_device = utilities_lib::device(args.force_cpu).unwrap();
-    let is_v3_model = args.which.is_v3_model();
-    let is_metal = !initial_device.is_cpu()
-        && candle_core::utils::metal_is_available()
-        && !args.force_cpu;
-
-    let device = if is_v3_model && is_metal {
-        candle_core::Device::Cpu
-    } else {
-        initial_device
-    };
-
-    let dtype = if device.is_cuda() { DType::BF16 } else { DType::F32 };
-    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&weight_paths, dtype, &device).unwrap() };
-
-    let model = match args.which {
-        Which::Base2B
-        | Which::Base7B
-        | Which::Instruct2B
-        | Which::Instruct7B
-        | Which::InstructV1_1_2B
-        | Which::InstructV1_1_7B
-        | Which::CodeBase2B
-        | Which::CodeBase7B
-        | Which::CodeInstruct2B
-        | Which::CodeInstruct7B => {
-            let config: Config1 = serde_json::from_reader(std::fs::File::open(config_path.clone()).unwrap()).unwrap();
-            GemmaModel::V1(Model1::new(args.use_flash_attn, &config, vb).unwrap())
-        }
-        Which::BaseV2_2B | Which::InstructV2_2B | Which::BaseV2_9B | Which::InstructV2_9B => {
-            let config: Config2 = serde_json::from_reader(std::fs::File::open(config_path.clone()).unwrap()).unwrap();
-            GemmaModel::V2(Model2::new(args.use_flash_attn, &config, vb).unwrap())
-        }
-        Which::BaseV3_1B | Which::InstructV3_1B => {
-            let config: Config3 = serde_json::from_reader(std::fs::File::open(config_path).unwrap()).unwrap();
-            GemmaModel::V3(Model3::new(args.use_flash_attn, &config, vb).unwrap())
-        }
-    };
-
-    TextGeneration::new(
-        model,
-        tokenizer,
-        args.seed,
-        args.temperature,
-        args.top_p,
-        args.repeat_penalty,
-        args.repeat_last_n,
-        &device,
-    )
+    model_id.to_lowercase().replace("_", "-")
 }
 
 fn build_gemma_prompt(messages: &[Message]) -> String {
     let mut prompt = String::new();
-    let mut system_prompt: Option<String> = None;
 
     for message in messages {
-        let content = match &message.content {
-            Some(content) => match &content.0 {
-                Either::Left(text) => text.clone(),
-                Either::Right(_) => "".to_string(),
-            },
-            None => "".to_string(),
-        };
-
         match message.role.as_str() {
-            "system" => system_prompt = Some(content),
-            "user" => {
-                prompt.push_str("<start_of_turn>user\n");
-                if let Some(sys_prompt) = system_prompt.take() {
-                    prompt.push_str(&sys_prompt);
-                    prompt.push_str("\n\n");
+            "system" => {
+                if let Some(MessageContent(Either::Left(content))) = &message.content {
+                    prompt.push_str(&format!(
+                        "<start_of_turn>system\n{}<end_of_turn>\n",
+                        content
+                    ));
+                }
+            }
+            "user" => {
+                if let Some(MessageContent(Either::Left(content))) = &message.content {
+                    prompt.push_str(&format!("<start_of_turn>user\n{}<end_of_turn>\n", content));
                 }
-                prompt.push_str(&content);
-                prompt.push_str("<end_of_turn>\n");
             }
             "assistant" => {
-                prompt.push_str("<start_of_turn>model\n");
-                prompt.push_str(&content);
-                prompt.push_str("<end_of_turn>\n");
+                if let Some(MessageContent(Either::Left(content))) = &message.content {
+                    prompt.push_str(&format!("<start_of_turn>model\n{}<end_of_turn>\n", content));
+                }
             }
             _ => {}
         }
@@ -316,23 +103,26 @@ pub async fn chat_completions(
     Json(request): Json<ChatCompletionRequest>,
 ) -> Result<impl IntoResponse, (StatusCode, String)> {
     if !request.stream.unwrap_or(false) {
-        return Ok(chat_completions_non_streaming_proxy(state, request).await.into_response());
+        return Ok(chat_completions_non_streaming_proxy(state, request)
+            .await
+            .into_response());
     }
-    Ok(chat_completions_stream(state, request).await.into_response())
+    Ok(chat_completions_stream(state, request)
+        .await
+        .into_response())
 }
 
 pub async fn chat_completions_non_streaming_proxy(
     state: AppState,
     request: ChatCompletionRequest,
 ) -> Result<impl IntoResponse, (StatusCode, Json<Value>)> {
-    let prompt = build_gemma_prompt(&request.messages);
-
     // Enforce model selection behavior: reject if a different model is requested
-    let configured_model = state.build_args.model_id.clone();
+    let configured_model = state.model_id.clone();
     let requested_model = request.model.clone();
     if requested_model.to_lowercase() != "default" {
         let normalized_requested = normalize_model_id(&requested_model);
-        if normalized_requested != configured_model {
+        let normalized_configured = normalize_model_id(&configured_model);
+        if normalized_requested != normalized_configured {
             return Err((
                 StatusCode::BAD_REQUEST,
                 Json(serde_json::json!({
@@ -349,35 +139,84 @@ pub async fn chat_completions_non_streaming_proxy(
     }
 
     let model_id = state.model_id.clone();
+    let max_tokens = request.max_tokens.unwrap_or(1000);
 
-    let mut buffer = Vec::new();
-    {
-        let mut text_gen = state.text_generation.lock().await;
-        // Reset per-request state without rebuilding the whole pipeline
-        text_gen.reset_state();
-        let max_tokens = request.max_tokens.unwrap_or(1000);
-        if let Err(e) = text_gen.run_with_output(&prompt, max_tokens, &mut buffer) {
-            return Err((
-                StatusCode::BAD_REQUEST,
-                Json(serde_json::json!({
-                    "error": { "message": format!("Error generating text: {}", e) }
-                })),
-            ));
-        }
-    }
-
-    let completion = match String::from_utf8(buffer) {
-        Ok(s) => s,
-        Err(e) => {
-            return Err((
-                StatusCode::INTERNAL_SERVER_ERROR,
-                Json(serde_json::json!({
-                    "error": { "message": format!("UTF-8 conversion error: {}", e) }
-                })),
-            ));
+    // Build prompt based on model type
+    let prompt = match state.model_type {
+        ModelType::Gemma => build_gemma_prompt(&request.messages),
+        ModelType::Llama => {
+            // For Llama, just use the last user message for now
+            request
+                .messages
+                .last()
+                .and_then(|m| m.content.as_ref())
+                .and_then(|c| match c {
+                    MessageContent(Either::Left(text)) => Some(text.clone()),
+                    _ => None,
+                })
+                .unwrap_or_default()
         }
     };
 
+    // Get streaming receiver based on model type
+    let rx =
+        match state.model_type {
+            ModelType::Gemma => {
+                if let Some(mut config) = state.gemma_config {
+                    config.prompt = prompt.clone();
+                    config.max_tokens = max_tokens;
+                    run_gemma_api(config).map_err(|e| (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    Json(serde_json::json!({
+                        "error": { "message": format!("Error initializing Gemma model: {}", e) }
+                    }))
+                ))?
+                } else {
+                    return Err((
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        Json(serde_json::json!({
+                            "error": { "message": "Gemma configuration not available" }
+                        })),
+                    ));
+                }
+            }
+            ModelType::Llama => {
+                if let Some(mut config) = state.llama_config {
+                    config.prompt = prompt.clone();
+                    config.max_tokens = max_tokens;
+                    run_llama_inference(config).map_err(|e| (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    Json(serde_json::json!({
+                        "error": { "message": format!("Error initializing Llama model: {}", e) }
+                    }))
+                ))?
+                } else {
+                    return Err((
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        Json(serde_json::json!({
+                            "error": { "message": "Llama configuration not available" }
+                        })),
+                    ));
+                }
+            }
+        };
+
+    // Collect all tokens from the stream
+    let mut completion = String::new();
+    while let Ok(token_result) = rx.recv() {
+        match token_result {
+            Ok(token) => completion.push_str(&token),
+            Err(e) => {
+                return Err((
+                    StatusCode::BAD_REQUEST,
+                    Json(serde_json::json!({
+                        "error": { "message": format!("Error generating text: {}", e) }
+                    })),
+                ));
+            }
+        }
+    }
+
     let response = ChatCompletionResponse {
         id: format!("chatcmpl-{}", Uuid::new_v4().to_string().replace('-', "")),
         object: "chat.completion".to_string(),
@@ -420,11 +259,12 @@ async fn handle_streaming_request(
     request: ChatCompletionRequest,
 ) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (StatusCode, Json<Value>)> {
     // Validate requested model vs configured model
-    let configured_model = state.build_args.model_id.clone();
+    let configured_model = state.model_id.clone();
     let requested_model = request.model.clone();
     if requested_model.to_lowercase() != "default" {
         let normalized_requested = normalize_model_id(&requested_model);
-        if normalized_requested != configured_model {
+        let normalized_configured = normalize_model_id(&configured_model);
+        if normalized_requested != normalized_configured {
             return Err((
                 StatusCode::BAD_REQUEST,
                 Json(serde_json::json!({
@@ -447,9 +287,24 @@ async fn handle_streaming_request(
         .unwrap_or_default()
         .as_secs();
     let model_id = state.model_id.clone();
+    let max_tokens = request.max_tokens.unwrap_or(1000);
 
-    // Build prompt
-    let prompt = build_gemma_prompt(&request.messages);
+    // Build prompt based on model type
+    let prompt = match state.model_type {
+        ModelType::Gemma => build_gemma_prompt(&request.messages),
+        ModelType::Llama => {
+            // For Llama, just use the last user message for now
+            request
+                .messages
+                .last()
+                .and_then(|m| m.content.as_ref())
+                .and_then(|c| match c {
+                    MessageContent(Either::Left(text)) => Some(text.clone()),
+                    _ => None,
+                })
+                .unwrap_or_default()
+        }
+    };
     tracing::debug!("Formatted prompt: {}", prompt);
 
     // Channel for streaming SSE events
@@ -463,7 +318,10 @@ async fn handle_streaming_request(
         model: model_id.clone(),
         choices: vec![ChatCompletionChunkChoice {
             index: 0,
-            delta: Delta { role: Some("assistant".to_string()), content: None },
+            delta: Delta {
+                role: Some("assistant".to_string()),
+                content: None,
+            },
             finish_reason: None,
         }],
     };
@@ -471,80 +329,128 @@ async fn handle_streaming_request(
         let _ = tx.send(Ok(Event::default().data(json)));
     }
 
-    // Spawn generation task that streams tokens as they are generated
-    let state_clone = state.clone();
-    let response_id_clone = response_id.clone();
-    tokio::spawn(async move {
-        let max_tokens = request.max_tokens.unwrap_or(1000);
-        let mut text_gen = state_clone.text_generation.lock().await;
-        text_gen.reset_state();
+    // Get streaming receiver based on model type
+    let model_rx = match state.model_type {
+        ModelType::Gemma => {
+            if let Some(mut config) = state.gemma_config {
+                config.prompt = prompt.clone();
+                config.max_tokens = max_tokens;
+                match run_gemma_api(config) {
+                    Ok(rx) => rx,
+                    Err(e) => {
+                        return Err((
+                            StatusCode::INTERNAL_SERVER_ERROR,
+                            Json(serde_json::json!({
+                                "error": { "message": format!("Error initializing Gemma model: {}", e) }
+                            })),
+                        ));
+                    }
+                }
+            } else {
+                return Err((
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    Json(serde_json::json!({
+                        "error": { "message": "Gemma configuration not available" }
+                    })),
+                ));
+            }
+        }
+        ModelType::Llama => {
+            if let Some(mut config) = state.llama_config {
+                config.prompt = prompt.clone();
+                config.max_tokens = max_tokens;
+                match run_llama_inference(config) {
+                    Ok(rx) => rx,
+                    Err(e) => {
+                        return Err((
+                            StatusCode::INTERNAL_SERVER_ERROR,
+                            Json(serde_json::json!({
+                                "error": { "message": format!("Error initializing Llama model: {}", e) }
+                            })),
+                        ));
+                    }
+                }
+            } else {
+                return Err((
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    Json(serde_json::json!({
+                        "error": { "message": "Llama configuration not available" }
+                    })),
+                ));
+            }
+        }
+    };
 
-        // Stream tokens via callback with repetition detection
+    // Spawn task to receive tokens from model and forward as SSE events
+    let response_id_clone = response_id.clone();
+    let model_id_clone = model_id.clone();
+    tokio::spawn(async move {
+        // Stream tokens with repetition detection
         let mut recent_tokens = Vec::new();
         let mut repetition_count = 0;
-        const MAX_REPETITION_COUNT: usize = 5; // Stop after 5 consecutive repetitions
-        const REPETITION_WINDOW: usize = 8; // Look at last 8 tokens for patterns
-        
-        let result = text_gen.run_with_streaming(&prompt, max_tokens, |token| {
-            // Debug log to verify token content
-            tracing::debug!("Streaming token: '{}'", token);
-            
-            // Skip sending empty tokens
-            if token.is_empty() {
-                tracing::debug!("Skipping empty token");
-                return Ok(());
-            }
-            
-            // Add token to recent history for repetition detection
-            recent_tokens.push(token.to_string());
-            if recent_tokens.len() > REPETITION_WINDOW {
-                recent_tokens.remove(0);
-            }
-            
-            // Check for repetitive patterns
-            if recent_tokens.len() >= 4 {
-                let last_token = &recent_tokens[recent_tokens.len() - 1];
-                let second_last = &recent_tokens[recent_tokens.len() - 2];
-                
-                // Check if we're repeating the same token or pattern
-                if last_token == second_last || 
-                   (last_token.trim() == "plus" && second_last.trim() == "plus") ||
-                   (recent_tokens.len() >= 6 && 
-                    recent_tokens[recent_tokens.len()-3..].iter().all(|t| t.trim() == "plus" || t.trim().is_empty())) {
-                    repetition_count += 1;
-                    tracing::warn!("Detected repetition pattern: '{}' (count: {})", last_token, repetition_count);
-                    
-                    if repetition_count >= MAX_REPETITION_COUNT {
-                        tracing::info!("Stopping generation due to excessive repetition");
-                        return Err(anyhow::Error::msg("Repetition detected - stopping generation"));
+        const MAX_REPETITION_COUNT: usize = 5;
+        const REPETITION_WINDOW: usize = 8;
+
+        while let Ok(token_result) = model_rx.recv() {
+            match token_result {
+                Ok(token) => {
+                    // Skip sending empty tokens
+                    if token.is_empty() {
+                        continue;
                     }
-                } else {
-                    repetition_count = 0; // Reset counter if pattern breaks
+
+                    // Add token to recent history for repetition detection
+                    recent_tokens.push(token.clone());
+                    if recent_tokens.len() > REPETITION_WINDOW {
+                        recent_tokens.remove(0);
+                    }
+
+                    // Check for repetitive patterns
+                    if recent_tokens.len() >= 4 {
+                        let last_token = &recent_tokens[recent_tokens.len() - 1];
+                        let second_last = &recent_tokens[recent_tokens.len() - 2];
+
+                        if last_token == second_last {
+                            repetition_count += 1;
+                            tracing::warn!(
+                                "Detected repetition pattern: '{}' (count: {})",
+                                last_token,
+                                repetition_count
+                            );
+
+                            if repetition_count >= MAX_REPETITION_COUNT {
+                                tracing::info!("Stopping generation due to excessive repetition");
+                                break;
+                            }
+                        } else {
+                            repetition_count = 0;
+                        }
+                    }
+
+                    let chunk = ChatCompletionChunk {
+                        id: response_id_clone.clone(),
+                        object: "chat.completion.chunk".to_string(),
+                        created,
+                        model: model_id_clone.clone(),
+                        choices: vec![ChatCompletionChunkChoice {
+                            index: 0,
+                            delta: Delta {
+                                role: None,
+                                content: Some(token),
+                            },
+                            finish_reason: None,
+                        }],
+                    };
+
+                    if let Ok(json) = serde_json::to_string(&chunk) {
+                        let _ = tx.send(Ok(Event::default().data(json)));
+                    }
+                }
+                Err(e) => {
+                    tracing::info!("Text generation stopped: {}", e);
+                    break;
                 }
             }
-            
-            let chunk = ChatCompletionChunk {
-                id: response_id_clone.clone(),
-                object: "chat.completion.chunk".to_string(),
-                created,
-                model: model_id.clone(),
-                choices: vec![ChatCompletionChunkChoice {
-                    index: 0,
-                    delta: Delta { role: None, content: Some(token.to_string()) },
-                    finish_reason: None,
-                }],
-            };
-            if let Ok(json) = serde_json::to_string(&chunk) {
-                tracing::debug!("Sending chunk with content: '{}'", token);
-                let _ = tx.send(Ok(Event::default().data(json)));
-            }
-            Ok(())
-        }).await;
-        
-        // Log result of generation
-        match result {
-            Ok(_) => tracing::debug!("Text generation completed successfully"),
-            Err(e) => tracing::info!("Text generation stopped: {}", e),
         }
 
         // Send final stop chunk and DONE marker
@@ -552,10 +458,13 @@ async fn handle_streaming_request(
             id: response_id_clone.clone(),
             object: "chat.completion.chunk".to_string(),
             created,
-            model: model_id.clone(),
+            model: model_id_clone.clone(),
             choices: vec![ChatCompletionChunkChoice {
                 index: 0,
-                delta: Delta { role: None, content: None },
+                delta: Delta {
+                    role: None,
+                    content: None,
+                },
                 finish_reason: Some("stop".to_string()),
             }],
         };
@@ -570,8 +479,6 @@ async fn handle_streaming_request(
     Ok(Sse::new(stream))
 }
 
-
-
 // -------------------------
 // Router
 // -------------------------
@@ -594,6 +501,7 @@ pub fn create_router(app_state: AppState) -> Router {
 pub async fn list_models() -> Json<ModelListResponse> {
     // Get all available model variants from the Which enum
     let models = vec![
+        // Gemma models
         Model {
             id: "gemma-2b".to_string(),
             object: "model".to_string(),
@@ -690,6 +598,73 @@ pub async fn list_models() -> Json<ModelListResponse> {
             created: 1686935002,
             owned_by: "google".to_string(),
         },
+        // Llama models
+        Model {
+            id: "llama-3.2-1b".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "meta".to_string(),
+        },
+        Model {
+            id: "llama-3.2-1b-instruct".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "meta".to_string(),
+        },
+        Model {
+            id: "llama-3.2-3b".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "meta".to_string(),
+        },
+        Model {
+            id: "llama-3.2-3b-instruct".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "meta".to_string(),
+        },
+        Model {
+            id: "smollm2-135m".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "huggingface".to_string(),
+        },
+        Model {
+            id: "smollm2-135m-instruct".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "huggingface".to_string(),
+        },
+        Model {
+            id: "smollm2-360m".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "huggingface".to_string(),
+        },
+        Model {
+            id: "smollm2-360m-instruct".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "huggingface".to_string(),
+        },
+        Model {
+            id: "smollm2-1.7b".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "huggingface".to_string(),
+        },
+        Model {
+            id: "smollm2-1.7b-instruct".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "huggingface".to_string(),
+        },
+        Model {
+            id: "tinyllama-1.1b-chat".to_string(),
+            object: "model".to_string(),
+            created: 1686935002,
+            owned_by: "tinyllama".to_string(),
+        },
     ];
 
     Json(ModelListResponse {
@@ -698,7 +673,6 @@ pub async fn list_models() -> Json<ModelListResponse> {
     })
 }
 
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -732,10 +706,7 @@ mod tests {
 
         let prompt = build_gemma_prompt(&messages);
 
-        let expected = "<start_of_turn>user\nSystem message\n\nKnock knock.<end_of_turn>\n\
-                       <start_of_turn>model\nWho's there?<end_of_turn>\n\
-                       <start_of_turn>user\nGemma.<end_of_turn>\n\
-                       <start_of_turn>model\n";
+        let expected = "<start_of_turn>system\nSystem message<end_of_turn>\n<start_of_turn>user\nKnock knock.<end_of_turn>\n<start_of_turn>model\nWho's there?<end_of_turn>\n<start_of_turn>user\nGemma.<end_of_turn>\n<start_of_turn>model\n";
 
         assert_eq!(prompt, expected);
     }
@@ -749,15 +720,13 @@ mod tests {
 
     #[test]
     fn test_missing_content() {
-        let messages = vec![
-            Message {
-                role: "user".to_string(),
-                content: None,
-                name: None,
-            }
-        ];
+        let messages = vec![Message {
+            role: "user".to_string(),
+            content: None,
+            name: None,
+        }];
 
         let prompt = build_gemma_prompt(&messages);
-        assert_eq!(prompt, "<start_of_turn>user\n<end_of_turn>\n<start_of_turn>model\n");
+        assert_eq!(prompt, "<start_of_turn>model\n");
     }
 }

crates/inference-engine/src/token_output_stream.rs

@@ -1,87 +0,0 @@
-use candle_core::Result;
-
-/// This is a wrapper around a tokenizer to ensure that tokens can be returned to the user in a
-/// streaming way rather than having to wait for the full decoding.
-pub struct TokenOutputStream {
-    tokenizer: tokenizers::Tokenizer,
-    tokens: Vec<u32>,
-    prev_index: usize,
-    current_index: usize,
-}
-
-impl TokenOutputStream {
-    pub fn new(tokenizer: tokenizers::Tokenizer) -> Self {
-        Self {
-            tokenizer,
-            tokens: Vec::new(),
-            prev_index: 0,
-            current_index: 0,
-        }
-    }
-
-    pub fn into_inner(self) -> tokenizers::Tokenizer {
-        self.tokenizer
-    }
-
-    fn decode(&self, tokens: &[u32]) -> Result<String> {
-        match self.tokenizer.decode(tokens, true) {
-            Ok(str) => Ok(str),
-            Err(err) => candle_core::bail!("cannot decode: {err}"),
-        }
-    }
-
-    // https://github.com/huggingface/text-generation-inference/blob/5ba53d44a18983a4de32d122f4cb46f4a17d9ef6/server/text_generation_server/models/model.py#L68
-    pub fn next_token(&mut self, token: u32) -> Result<Option<String>> {
-        let prev_text = if self.tokens.is_empty() {
-            String::new()
-        } else {
-            let tokens = &self.tokens[self.prev_index..self.current_index];
-            self.decode(tokens)?
-        };
-        self.tokens.push(token);
-        let text = self.decode(&self.tokens[self.prev_index..])?;
-        if text.len() > prev_text.len() {
-            // Modified to include all tokens, not just alphanumeric ones
-            let text = text.split_at(prev_text.len());
-            self.prev_index = self.current_index;
-            self.current_index = self.tokens.len();
-            Ok(Some(text.1.to_string()))
-        } else {
-            Ok(None)
-        }
-    }
-
-    pub fn decode_rest(&self) -> Result<Option<String>> {
-        let prev_text = if self.tokens.is_empty() {
-            String::new()
-        } else {
-            let tokens = &self.tokens[self.prev_index..self.current_index];
-            self.decode(tokens)?
-        };
-        let text = self.decode(&self.tokens[self.prev_index..])?;
-        if text.len() > prev_text.len() {
-            let text = text.split_at(prev_text.len());
-            Ok(Some(text.1.to_string()))
-        } else {
-            Ok(None)
-        }
-    }
-
-    pub fn decode_all(&self) -> Result<String> {
-        self.decode(&self.tokens)
-    }
-
-    pub fn get_token(&self, token_s: &str) -> Option<u32> {
-        self.tokenizer.get_vocab(true).get(token_s).copied()
-    }
-
-    pub fn tokenizer(&self) -> &tokenizers::Tokenizer {
-        &self.tokenizer
-    }
-
-    pub fn clear(&mut self) {
-        self.tokens.clear();
-        self.prev_index = 0;
-        self.current_index = 0;
-    }
-}

crates/inference-engine/src/utilities_lib.rs

@@ -1,167 +0,0 @@
-use candle_core::utils::{cuda_is_available, metal_is_available};
-use candle_core::{Device, Result, Tensor};
-
-pub fn device(cpu: bool) -> Result<Device> {
-    if cpu {
-        Ok(Device::Cpu)
-    } else if cuda_is_available() {
-        Ok(Device::new_cuda(0)?)
-    } else if metal_is_available() {
-        Ok(Device::new_metal(0)?)
-    } else {
-        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
-        {
-            println!(
-                "Running on CPU, to run on GPU(metal), build this example with `--features metal`"
-            );
-        }
-        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
-        {
-            println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
-        }
-        Ok(Device::Cpu)
-    }
-}
-
-pub fn load_image<P: AsRef<std::path::Path>>(
-    p: P,
-    resize_longest: Option<usize>,
-) -> Result<(Tensor, usize, usize)> {
-    let img = image::ImageReader::open(p)?
-        .decode()
-        .map_err(candle_core::Error::wrap)?;
-    let (initial_h, initial_w) = (img.height() as usize, img.width() as usize);
-    let img = match resize_longest {
-        None => img,
-        Some(resize_longest) => {
-            let (height, width) = (img.height(), img.width());
-            let resize_longest = resize_longest as u32;
-            let (height, width) = if height < width {
-                let h = (resize_longest * height) / width;
-                (h, resize_longest)
-            } else {
-                let w = (resize_longest * width) / height;
-                (resize_longest, w)
-            };
-            img.resize_exact(width, height, image::imageops::FilterType::CatmullRom)
-        }
-    };
-    let (height, width) = (img.height() as usize, img.width() as usize);
-    let img = img.to_rgb8();
-    let data = img.into_raw();
-    let data = Tensor::from_vec(data, (height, width, 3), &Device::Cpu)?.permute((2, 0, 1))?;
-    Ok((data, initial_h, initial_w))
-}
-
-pub fn load_image_and_resize<P: AsRef<std::path::Path>>(
-    p: P,
-    width: usize,
-    height: usize,
-) -> Result<Tensor> {
-    let img = image::ImageReader::open(p)?
-        .decode()
-        .map_err(candle_core::Error::wrap)?
-        .resize_to_fill(
-            width as u32,
-            height as u32,
-            image::imageops::FilterType::Triangle,
-        );
-    let img = img.to_rgb8();
-    let data = img.into_raw();
-    Tensor::from_vec(data, (width, height, 3), &Device::Cpu)?.permute((2, 0, 1))
-}
-
-/// Saves an image to disk using the image crate, this expects an input with shape
-/// (c, height, width).
-pub fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
-    let p = p.as_ref();
-    let (channel, height, width) = img.dims3()?;
-    if channel != 3 {
-        candle_core::bail!("save_image expects an input of shape (3, height, width)")
-    }
-    let img = img.permute((1, 2, 0))?.flatten_all()?;
-    let pixels = img.to_vec1::<u8>()?;
-    let image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> =
-        match image::ImageBuffer::from_raw(width as u32, height as u32, pixels) {
-            Some(image) => image,
-            None => candle_core::bail!("error saving image {p:?}"),
-        };
-    image.save(p).map_err(candle_core::Error::wrap)?;
-    Ok(())
-}
-
-pub fn save_image_resize<P: AsRef<std::path::Path>>(
-    img: &Tensor,
-    p: P,
-    h: usize,
-    w: usize,
-) -> Result<()> {
-    let p = p.as_ref();
-    let (channel, height, width) = img.dims3()?;
-    if channel != 3 {
-        candle_core::bail!("save_image expects an input of shape (3, height, width)")
-    }
-    let img = img.permute((1, 2, 0))?.flatten_all()?;
-    let pixels = img.to_vec1::<u8>()?;
-    let image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> =
-        match image::ImageBuffer::from_raw(width as u32, height as u32, pixels) {
-            Some(image) => image,
-            None => candle_core::bail!("error saving image {p:?}"),
-        };
-    let image = image::DynamicImage::from(image);
-    let image = image.resize_to_fill(w as u32, h as u32, image::imageops::FilterType::CatmullRom);
-    image.save(p).map_err(candle_core::Error::wrap)?;
-    Ok(())
-}
-
-/// Loads the safetensors files for a model from the hub based on a json index file.
-pub fn hub_load_safetensors(
-    repo: &hf_hub::api::sync::ApiRepo,
-    json_file: &str,
-) -> Result<Vec<std::path::PathBuf>> {
-    let json_file = repo.get(json_file).map_err(candle_core::Error::wrap)?;
-    let json_file = std::fs::File::open(json_file)?;
-    let json: serde_json::Value =
-        serde_json::from_reader(&json_file).map_err(candle_core::Error::wrap)?;
-    let weight_map = match json.get("weight_map") {
-        None => candle_core::bail!("no weight map in {json_file:?}"),
-        Some(serde_json::Value::Object(map)) => map,
-        Some(_) => candle_core::bail!("weight map in {json_file:?} is not a map"),
-    };
-    let mut safetensors_files = std::collections::HashSet::new();
-    for value in weight_map.values() {
-        if let Some(file) = value.as_str() {
-            safetensors_files.insert(file.to_string());
-        }
-    }
-    let safetensors_files = safetensors_files
-        .iter()
-        .map(|v| repo.get(v).map_err(candle_core::Error::wrap))
-        .collect::<Result<Vec<_>>>()?;
-    Ok(safetensors_files)
-}
-
-pub fn hub_load_local_safetensors<P: AsRef<std::path::Path>>(
-    path: P,
-    json_file: &str,
-) -> Result<Vec<std::path::PathBuf>> {
-    let path = path.as_ref();
-    let jsfile = std::fs::File::open(path.join(json_file))?;
-    let json: serde_json::Value = serde_json::from_reader(&jsfile).map_err(candle_core::Error::wrap)?;
-    let weight_map = match json.get("weight_map") {
-        None => candle_core::bail!("no weight map in {json_file:?}"),
-        Some(serde_json::Value::Object(map)) => map,
-        Some(_) => candle_core::bail!("weight map in {json_file:?} is not a map"),
-    };
-    let mut safetensors_files = std::collections::HashSet::new();
-    for value in weight_map.values() {
-        if let Some(file) = value.as_str() {
-            safetensors_files.insert(file);
-        }
-    }
-    let safetensors_files: Vec<_> = safetensors_files
-        .into_iter()
-        .map(|v| path.join(v))
-        .collect();
-    Ok(safetensors_files)
-}

crates/inference-engine/tests/model_tests.rs

@@ -9,7 +9,10 @@ mod tests {
         // Test a few representative model variants
         assert_eq!(Which::Base2B.to_model_id(), "google/gemma-2b");
         assert_eq!(Which::Instruct7B.to_model_id(), "google/gemma-7b-it");
-        assert_eq!(Which::InstructV1_1_2B.to_model_id(), "google/gemma-1.1-2b-it");
+        assert_eq!(
+            Which::InstructV1_1_2B.to_model_id(),
+            "google/gemma-1.1-2b-it"
+        );
         assert_eq!(Which::CodeBase2B.to_model_id(), "google/codegemma-2b");
         assert_eq!(Which::BaseV2_2B.to_model_id(), "google/gemma-2-2b");
         assert_eq!(Which::InstructV3_1B.to_model_id(), "google/gemma-3-1b-it");
@@ -64,4 +67,4 @@ mod tests {
     // Note: Testing the Model enum's forward method would require creating actual model instances,
     // which is complex and would require loading model weights. This is better suited for
     // integration tests or mocking the models.
-}
+}

crates/inference-engine/tests/text_generation_tests.rs

@@ -1,549 +0,0 @@
-use anyhow::Result;
-use candle_core::{Device, Tensor};
-use candle_transformers::generation::LogitsProcessor;
-use inference_engine::model::Which;
-use inference_engine::text_generation::TextGeneration;
-use inference_engine::token_output_stream::TokenOutputStream;
-use std::collections::HashMap;
-use tokenizers::Tokenizer;
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    // Helper function to create a simple tokenizer for testing
-    fn create_test_tokenizer() -> Result<Tokenizer> {
-        // Create a simple tokenizer from the pretrained model
-        // This uses the tokenizer from the Hugging Face hub
-        let tokenizer = Tokenizer::from_pretrained("google/gemma-2b", None).unwrap();
-        Ok(tokenizer)
-    }
-
-    // Test the Which enum's to_model_id method
-    #[test]
-    fn test_which_model_id() {
-        assert_eq!(Which::Base2B.to_model_id(), "google/gemma-2b");
-        assert_eq!(Which::Instruct7B.to_model_id(), "google/gemma-7b-it");
-    }
-
-    // Test the Which enum's is_instruct_model method
-    #[test]
-    fn test_which_is_instruct() {
-        assert!(!Which::Base2B.is_instruct_model());
-        assert!(Which::Instruct7B.is_instruct_model());
-    }
-
-    // Test the Which enum's is_v3_model method
-    #[test]
-    fn test_which_is_v3() {
-        assert!(!Which::Base2B.is_v3_model());
-        assert!(Which::BaseV3_1B.is_v3_model());
-    }
-
-    // Test the TokenOutputStream functionality
-    #[test]
-    fn test_token_output_stream() -> Result<()> {
-        let tokenizer = create_test_tokenizer()?;
-        let mut token_stream = TokenOutputStream::new(tokenizer);
-
-        // Test encoding and decoding
-        let text = "Hello, world!";
-        let encoded = token_stream.tokenizer().encode(text, true).unwrap();
-        let token_ids = encoded.get_ids();
-
-        // Add tokens one by one
-        for &token_id in token_ids {
-            token_stream.next_token(token_id)?;
-        }
-
-        // Decode all and check
-        let decoded = token_stream.decode_all()?;
-        assert_eq!(decoded.trim(), text);
-
-        Ok(())
-    }
-
-    // Test the LogitsProcessor
-    #[test]
-    fn test_logits_processor() -> Result<()> {
-        // Create a LogitsProcessor with default settings
-        let seed = 42;
-        let temp = Some(0.8);
-        let top_p = Some(0.9);
-        let logits_processor = LogitsProcessor::new(seed, temp, top_p);
-
-        // Create a simple logits tensor
-        // In a real test, we would create a tensor with known values and verify
-        // that sampling produces expected results
-
-        // For now, we'll just verify that the LogitsProcessor can be created
-        assert!(true);
-        Ok(())
-    }
-
-    // Test the TextGeneration constructor
-    #[test]
-    fn test_text_generation_constructor() -> Result<()> {
-        // We can't easily create a Model instance for testing,
-        // but we can test that the constructor compiles and the types are correct
-
-        // In a real test with a mock Model, we would:
-        // 1. Create a mock model
-        // 2. Create a tokenizer
-        // 3. Call TextGeneration::new
-        // 4. Verify the properties of the created instance
-
-        // For now, we'll just verify that the code compiles
-        assert!(true);
-        Ok(())
-    }
-
-    // Test apply_cached_repeat_penalty method with no penalty
-    #[test]
-    fn test_apply_cached_repeat_penalty_no_penalty() -> Result<()> {
-        // Create a simple test setup
-        let device = Device::Cpu;
-        let logits_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
-        let logits = Tensor::new(&logits_data[..], &device)?;
-        let tokens = vec![1u32, 2u32, 3u32];
-        
-        // Create a mock TextGeneration instance
-        // Since we can't easily create a full TextGeneration instance without a model,
-        // we'll test the logic by creating a simple struct with the necessary fields
-        struct MockTextGeneration {
-            repeat_penalty: f32,
-            repeat_last_n: usize,
-            penalty_cache: HashMap<usize, f32>,
-        }
-        
-        impl MockTextGeneration {
-            fn apply_cached_repeat_penalty(
-                &mut self,
-                logits: Tensor,
-                tokens: &[u32],
-            ) -> Result<(Tensor, std::time::Duration)> {
-                let repeat_start = std::time::Instant::now();
-
-                // If no penalty, return the original logits
-                if self.repeat_penalty == 1.0 {
-                    return Ok((logits, repeat_start.elapsed()));
-                }
-
-                // Get the tokens to penalize (the last n tokens)
-                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
-                let penalty_tokens = &tokens[start_at..];
-
-                // Extract logits to a vector for modification
-                let mut logits_vec = logits.to_vec1::<f32>()?;
-                let cache_hits = std::cell::Cell::new(0);
-
-                // Apply penalties with caching
-                for &token_id in penalty_tokens {
-                    let token_id = token_id as usize;
-                    if token_id < logits_vec.len() {
-                        // Check if we've already calculated this token's penalty
-                        if let Some(penalized_score) = self.penalty_cache.get(&token_id) {
-                            // Use cached value
-                            logits_vec[token_id] = *penalized_score;
-                            cache_hits.set(cache_hits.get() + 1);
-                        } else {
-                            // Calculate and cache new value
-                            let score = logits_vec[token_id];
-                            let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                            let penalized_score = sign * score / self.repeat_penalty;
-                            logits_vec[token_id] = penalized_score;
-                            self.penalty_cache.insert(token_id, penalized_score);
-                        }
-                    }
-                }
-
-                // Create a new tensor with the modified logits
-                let device = logits.device().clone();
-                let shape = logits.shape().clone();
-                let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                let result = new_logits.reshape(shape)?;
-
-                let elapsed = repeat_start.elapsed();
-                Ok((result, elapsed))
-            }
-        }
-        
-        let mut mock_gen = MockTextGeneration {
-            repeat_penalty: 1.0, // No penalty
-            repeat_last_n: 3,
-            penalty_cache: HashMap::new(),
-        };
-        
-        let (result_logits, _duration) = mock_gen.apply_cached_repeat_penalty(logits.clone(), &tokens)?;
-        let result_data = result_logits.to_vec1::<f32>()?;
-        
-        // With no penalty, logits should be unchanged
-        assert_eq!(result_data, logits_data);
-        Ok(())
-    }
-
-    // Test apply_cached_repeat_penalty method with penalty
-    #[test]
-    fn test_apply_cached_repeat_penalty_with_penalty() -> Result<()> {
-        let device = Device::Cpu;
-        let logits_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
-        let logits = Tensor::new(&logits_data[..], &device)?;
-        let tokens = vec![1u32, 2u32, 3u32];
-        
-        struct MockTextGeneration {
-            repeat_penalty: f32,
-            repeat_last_n: usize,
-            penalty_cache: HashMap<usize, f32>,
-        }
-        
-        impl MockTextGeneration {
-            fn apply_cached_repeat_penalty(
-                &mut self,
-                logits: Tensor,
-                tokens: &[u32],
-            ) -> Result<(Tensor, std::time::Duration)> {
-                let repeat_start = std::time::Instant::now();
-
-                if self.repeat_penalty == 1.0 {
-                    return Ok((logits, repeat_start.elapsed()));
-                }
-
-                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
-                let penalty_tokens = &tokens[start_at..];
-                let mut logits_vec = logits.to_vec1::<f32>()?;
-                let cache_hits = std::cell::Cell::new(0);
-
-                for &token_id in penalty_tokens {
-                    let token_id = token_id as usize;
-                    if token_id < logits_vec.len() {
-                        if let Some(penalized_score) = self.penalty_cache.get(&token_id) {
-                            logits_vec[token_id] = *penalized_score;
-                            cache_hits.set(cache_hits.get() + 1);
-                        } else {
-                            let score = logits_vec[token_id];
-                            let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                            let penalized_score = sign * score / self.repeat_penalty;
-                            logits_vec[token_id] = penalized_score;
-                            self.penalty_cache.insert(token_id, penalized_score);
-                        }
-                    }
-                }
-
-                let device = logits.device().clone();
-                let shape = logits.shape().clone();
-                let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                let result = new_logits.reshape(shape)?;
-
-                let elapsed = repeat_start.elapsed();
-                Ok((result, elapsed))
-            }
-        }
-        
-        let mut mock_gen = MockTextGeneration {
-            repeat_penalty: 2.0, // Apply penalty
-            repeat_last_n: 3,
-            penalty_cache: HashMap::new(),
-        };
-        
-        let (result_logits, _duration) = mock_gen.apply_cached_repeat_penalty(logits.clone(), &tokens)?;
-        let result_data = result_logits.to_vec1::<f32>()?;
-        
-        // Tokens 1, 2, 3 should be penalized (divided by 2.0)
-        let expected = vec![1.0f32, 1.0, 1.5, 2.0, 5.0]; // [1.0, 2.0/2.0, 3.0/2.0, 4.0/2.0, 5.0]
-        assert_eq!(result_data, expected);
-        Ok(())
-    }
-
-    // Test apply_cached_repeat_penalty caching behavior
-    #[test]
-    fn test_apply_cached_repeat_penalty_caching() -> Result<()> {
-        let device = Device::Cpu;
-        let logits_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
-        let logits = Tensor::new(&logits_data[..], &device)?;
-        let tokens = vec![1u32, 1u32, 1u32]; // Repeated token should use cache
-        
-        struct MockTextGeneration {
-            repeat_penalty: f32,
-            repeat_last_n: usize,
-            penalty_cache: HashMap<usize, f32>,
-        }
-        
-        impl MockTextGeneration {
-            fn apply_cached_repeat_penalty(
-                &mut self,
-                logits: Tensor,
-                tokens: &[u32],
-            ) -> Result<(Tensor, std::time::Duration)> {
-                let repeat_start = std::time::Instant::now();
-
-                if self.repeat_penalty == 1.0 {
-                    return Ok((logits, repeat_start.elapsed()));
-                }
-
-                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
-                let penalty_tokens = &tokens[start_at..];
-                let mut logits_vec = logits.to_vec1::<f32>()?;
-
-                for &token_id in penalty_tokens {
-                    let token_id = token_id as usize;
-                    if token_id < logits_vec.len() {
-                        if let Some(penalized_score) = self.penalty_cache.get(&token_id) {
-                            logits_vec[token_id] = *penalized_score;
-                        } else {
-                            let score = logits_vec[token_id];
-                            let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                            let penalized_score = sign * score / self.repeat_penalty;
-                            logits_vec[token_id] = penalized_score;
-                            self.penalty_cache.insert(token_id, penalized_score);
-                        }
-                    }
-                }
-
-                let device = logits.device().clone();
-                let shape = logits.shape().clone();
-                let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                let result = new_logits.reshape(shape)?;
-
-                let elapsed = repeat_start.elapsed();
-                Ok((result, elapsed))
-            }
-        }
-        
-        let mut mock_gen = MockTextGeneration {
-            repeat_penalty: 2.0,
-            repeat_last_n: 3,
-            penalty_cache: HashMap::new(),
-        };
-        
-        // First call should cache the penalty for token 1
-        let (_result_logits, _duration) = mock_gen.apply_cached_repeat_penalty(logits.clone(), &tokens)?;
-        
-        // Cache should contain the penalized value for token 1
-        assert!(mock_gen.penalty_cache.contains_key(&1));
-        assert_eq!(mock_gen.penalty_cache.get(&1), Some(&1.0)); // 2.0 / 2.0 = 1.0
-        
-        Ok(())
-    }
-
-    // Test edge case: empty tokens array
-    #[test]
-    fn test_apply_cached_repeat_penalty_empty_tokens() -> Result<()> {
-        let device = Device::Cpu;
-        let logits_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
-        let logits = Tensor::new(&logits_data[..], &device)?;
-        let tokens: Vec<u32> = vec![]; // Empty tokens
-        
-        struct MockTextGeneration {
-            repeat_penalty: f32,
-            repeat_last_n: usize,
-            penalty_cache: HashMap<usize, f32>,
-        }
-        
-        impl MockTextGeneration {
-            fn apply_cached_repeat_penalty(
-                &mut self,
-                logits: Tensor,
-                tokens: &[u32],
-            ) -> Result<(Tensor, std::time::Duration)> {
-                let repeat_start = std::time::Instant::now();
-
-                if self.repeat_penalty == 1.0 {
-                    return Ok((logits, repeat_start.elapsed()));
-                }
-
-                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
-                let penalty_tokens = &tokens[start_at..];
-                let mut logits_vec = logits.to_vec1::<f32>()?;
-
-                for &token_id in penalty_tokens {
-                    let token_id = token_id as usize;
-                    if token_id < logits_vec.len() {
-                        if let Some(penalized_score) = self.penalty_cache.get(&token_id) {
-                            logits_vec[token_id] = *penalized_score;
-                        } else {
-                            let score = logits_vec[token_id];
-                            let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                            let penalized_score = sign * score / self.repeat_penalty;
-                            logits_vec[token_id] = penalized_score;
-                            self.penalty_cache.insert(token_id, penalized_score);
-                        }
-                    }
-                }
-
-                let device = logits.device().clone();
-                let shape = logits.shape().clone();
-                let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                let result = new_logits.reshape(shape)?;
-
-                let elapsed = repeat_start.elapsed();
-                Ok((result, elapsed))
-            }
-        }
-        
-        let mut mock_gen = MockTextGeneration {
-            repeat_penalty: 2.0,
-            repeat_last_n: 3,
-            penalty_cache: HashMap::new(),
-        };
-        
-        let (result_logits, _duration) = mock_gen.apply_cached_repeat_penalty(logits.clone(), &tokens)?;
-        let result_data = result_logits.to_vec1::<f32>()?;
-        
-        // With empty tokens, logits should be unchanged
-        assert_eq!(result_data, logits_data);
-        Ok(())
-    }
-
-    // Test edge case: out-of-bounds token IDs
-    #[test]
-    fn test_apply_cached_repeat_penalty_out_of_bounds() -> Result<()> {
-        let device = Device::Cpu;
-        let logits_data = vec![1.0f32, 2.0, 3.0];
-        let logits = Tensor::new(&logits_data[..], &device)?;
-        let tokens = vec![1u32, 5u32, 10u32]; // Token 5 and 10 are out of bounds
-        
-        struct MockTextGeneration {
-            repeat_penalty: f32,
-            repeat_last_n: usize,
-            penalty_cache: HashMap<usize, f32>,
-        }
-        
-        impl MockTextGeneration {
-            fn apply_cached_repeat_penalty(
-                &mut self,
-                logits: Tensor,
-                tokens: &[u32],
-            ) -> Result<(Tensor, std::time::Duration)> {
-                let repeat_start = std::time::Instant::now();
-
-                if self.repeat_penalty == 1.0 {
-                    return Ok((logits, repeat_start.elapsed()));
-                }
-
-                let start_at = tokens.len().saturating_sub(self.repeat_last_n);
-                let penalty_tokens = &tokens[start_at..];
-                let mut logits_vec = logits.to_vec1::<f32>()?;
-
-                for &token_id in penalty_tokens {
-                    let token_id = token_id as usize;
-                    if token_id < logits_vec.len() {
-                        if let Some(penalized_score) = self.penalty_cache.get(&token_id) {
-                            logits_vec[token_id] = *penalized_score;
-                        } else {
-                            let score = logits_vec[token_id];
-                            let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                            let penalized_score = sign * score / self.repeat_penalty;
-                            logits_vec[token_id] = penalized_score;
-                            self.penalty_cache.insert(token_id, penalized_score);
-                        }
-                    }
-                }
-
-                let device = logits.device().clone();
-                let shape = logits.shape().clone();
-                let new_logits = Tensor::new(&logits_vec[..], &device)?;
-                let result = new_logits.reshape(shape)?;
-
-                let elapsed = repeat_start.elapsed();
-                Ok((result, elapsed))
-            }
-        }
-        
-        let mut mock_gen = MockTextGeneration {
-            repeat_penalty: 2.0,
-            repeat_last_n: 3,
-            penalty_cache: HashMap::new(),
-        };
-        
-        let (result_logits, _duration) = mock_gen.apply_cached_repeat_penalty(logits.clone(), &tokens)?;
-        let result_data = result_logits.to_vec1::<f32>()?;
-        
-        // Only token 1 should be penalized, out-of-bounds tokens should be ignored
-        let expected = vec![1.0f32, 1.0, 3.0]; // [1.0, 2.0/2.0, 3.0]
-        assert_eq!(result_data, expected);
-        Ok(())
-    }
-
-    // Test the actual apply_cached_repeat_penalty method from TextGeneration
-    // This test creates a TextGeneration instance with minimal dependencies to test the real method
-    #[test]
-    fn test_actual_apply_cached_repeat_penalty_implementation() -> Result<()> {
-        // Since creating a real TextGeneration instance requires a Model which needs model weights,
-        // we'll create a test that demonstrates the method is now public and can be accessed.
-        // The comprehensive functionality testing is already covered by the mock tests above.
-        
-        // Test data setup
-        let device = Device::Cpu;
-        let logits_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
-        let logits = Tensor::new(&logits_data[..], &device)?;
-        let tokens = vec![1u32, 2u32, 3u32];
-        
-        // Test that we can create the necessary components
-        let tokenizer = create_test_tokenizer()?;
-        
-        // The method is now public as confirmed by making it pub fn apply_cached_repeat_penalty
-        // This test verifies the method signature and that it's accessible from external code
-        
-        // We could create a TextGeneration instance if we had a way to mock the Model,
-        // but for now we confirm that the existing mock tests cover the functionality
-        // and the method is properly exposed as public
-        
-        println!("apply_cached_repeat_penalty method is now public and accessible for testing");
-        assert!(true);
-        Ok(())
-    }
-    
-    // Integration test that demonstrates the method usage pattern
-    #[test] 
-    fn test_apply_cached_repeat_penalty_usage_pattern() -> Result<()> {
-        // This test demonstrates how the apply_cached_repeat_penalty method would be used
-        // in practice, even though we can't create a full TextGeneration instance in unit tests
-        
-        let device = Device::Cpu;
-        let logits_data = vec![1.5f32, 2.5, 3.5, 4.5, 5.5];
-        let logits = Tensor::new(&logits_data[..], &device)?;
-        let tokens = vec![1u32, 2u32, 1u32, 3u32]; // Repeated token 1 to test caching
-        
-        // Test parameters that would be used with TextGeneration
-        let repeat_penalty = 1.2f32;
-        let repeat_last_n = 3usize;
-        let mut penalty_cache: HashMap<usize, f32> = HashMap::new();
-        
-        // Simulate the method's logic to verify it works as expected
-        let start_time = std::time::Instant::now();
-        
-        if repeat_penalty != 1.0 {
-            let start_at = tokens.len().saturating_sub(repeat_last_n);
-            let penalty_tokens = &tokens[start_at..];
-            let mut logits_vec = logits.to_vec1::<f32>()?;
-            
-            for &token_id in penalty_tokens {
-                let token_id = token_id as usize;
-                if token_id < logits_vec.len() {
-                    if let Some(_cached_score) = penalty_cache.get(&token_id) {
-                        // Cache hit simulation
-                    } else {
-                        let score = logits_vec[token_id];
-                        let sign = if score < 0.0 { -1.0 } else { 1.0 };
-                        let penalized_score = sign * score / repeat_penalty;
-                        penalty_cache.insert(token_id, penalized_score);
-                    }
-                }
-            }
-        }
-        
-        let _duration = start_time.elapsed();
-        
-        // Verify that tokens were processed correctly
-        assert!(penalty_cache.contains_key(&1)); // Token 1 should be cached
-        assert!(penalty_cache.contains_key(&2)); // Token 2 should be cached  
-        assert!(penalty_cache.contains_key(&3)); // Token 3 should be cached
-        
-        println!("Successfully demonstrated apply_cached_repeat_penalty usage pattern");
-        Ok(())
-    }
-
-    // Note: Testing the actual text generation functionality would require
-    // integration tests with real models, which is beyond the scope of these unit tests.
-    // The tests above focus on the components that can be tested in isolation.
-}

crates/inference-engine/tests/token_output_stream_tests.rs

@@ -1,129 +0,0 @@
-use inference_engine::token_output_stream::TokenOutputStream;
-use tokenizers::Tokenizer;
-use std::path::PathBuf;
-use anyhow::Result;
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    // Helper function to create a simple tokenizer for testing
-    fn create_test_tokenizer() -> Result<Tokenizer> {
-        // Create a simple tokenizer from the pretrained model
-        // This uses the tokenizer from the Hugging Face hub
-        let tokenizer = Tokenizer::from_pretrained("google/gemma-2b", None).unwrap();
-        Ok(tokenizer)
-    }
-
-    #[test]
-    fn test_new_token_output_stream() -> Result<()> {
-        let tokenizer = create_test_tokenizer()?;
-        let token_stream = TokenOutputStream::new(tokenizer);
-        
-        // Check that the token stream was created successfully
-        assert!(token_stream.tokenizer().get_vocab(true).len() > 0);
-        Ok(())
-    }
-
-    #[test]
-    fn test_clear() -> Result<()> {
-        let tokenizer = create_test_tokenizer()?;
-        let mut token_stream = TokenOutputStream::new(tokenizer);
-        
-        // Add a token
-        let token_id = token_stream.get_token("<eos>").unwrap();
-        token_stream.next_token(token_id)?;
-        
-        // Clear the stream
-        token_stream.clear();
-        
-        // Check that the stream is empty by trying to decode all
-        let decoded = token_stream.decode_all()?;
-        assert_eq!(decoded, "");
-        
-        Ok(())
-    }
-
-    #[test]
-    fn test_get_token() -> Result<()> {
-        let tokenizer = create_test_tokenizer()?;
-        let token_stream = TokenOutputStream::new(tokenizer);
-        
-        // Get a token that should exist
-        let eos_token = token_stream.get_token("<eos>");
-        assert!(eos_token.is_some());
-        
-        // Get a token that shouldn't exist
-        let nonexistent_token = token_stream.get_token("<this_token_does_not_exist>");
-        assert!(nonexistent_token.is_none());
-        
-        Ok(())
-    }
-
-    #[test]
-    fn test_next_token_and_decode() -> Result<()> {
-        let tokenizer = create_test_tokenizer()?;
-        let mut token_stream = TokenOutputStream::new(tokenizer);
-        
-        // Get some tokens
-        let hello_tokens = token_stream.tokenizer().encode("Hello world", true).unwrap();
-        let token_ids = hello_tokens.get_ids();
-        
-        // Add tokens one by one
-        let mut output = String::new();
-        for &token_id in token_ids {
-            if let Some(text) = token_stream.next_token(token_id)? {
-                output.push_str(&text);
-            }
-        }
-        
-        // Get any remaining text
-        if let Some(rest) = token_stream.decode_rest()? {
-            output.push_str(&rest);
-        }
-        
-        // Check the output
-        assert!(!output.is_empty());
-        assert_eq!(output.trim(), "Hello world");
-        
-        Ok(())
-    }
-
-    #[test]
-    fn test_decode_all() -> Result<()> {
-        let tokenizer = create_test_tokenizer()?;
-        let mut token_stream = TokenOutputStream::new(tokenizer);
-        
-        // Get some tokens
-        let hello_tokens = token_stream.tokenizer().encode("Hello world", true).unwrap();
-        let token_ids = hello_tokens.get_ids();
-        
-        // Add tokens one by one
-        for &token_id in token_ids {
-            token_stream.next_token(token_id)?;
-        }
-        
-        // Decode all
-        let decoded = token_stream.decode_all()?;
-        
-        // Check the output
-        assert_eq!(decoded.trim(), "Hello world");
-        
-        Ok(())
-    }
-
-    #[test]
-    fn test_into_inner() -> Result<()> {
-        let tokenizer = create_test_tokenizer()?;
-        let token_stream = TokenOutputStream::new(tokenizer);
-        
-        // Get the inner tokenizer
-        let inner_tokenizer = token_stream.into_inner();
-        
-        // Check that the inner tokenizer works
-        let encoded = inner_tokenizer.encode("Test", true).unwrap();
-        assert!(encoded.get_ids().len() > 0);
-        
-        Ok(())
-    }
-}

crates/leptos-app/src/app.rs

@@ -5,6 +5,25 @@ use leptos_router::{
     StaticSegment,
 };
 
+#[cfg(feature = "hydrate")]
+use async_openai_wasm::config::OpenAIConfig;
+#[cfg(feature = "hydrate")]
+use async_openai_wasm::types::{FinishReason, Role};
+#[cfg(feature = "hydrate")]
+use async_openai_wasm::{
+    types::{
+        ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestSystemMessageArgs,
+        ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs,
+        Model as OpenAIModel,
+    },
+    Client,
+};
+#[cfg(feature = "hydrate")]
+use futures_util::StreamExt;
+#[cfg(feature = "hydrate")]
+use js_sys::Date;
+#[cfg(feature = "hydrate")]
+use leptos::task::spawn_local;
 #[cfg(feature = "hydrate")]
 use serde::{Deserialize, Serialize};
 #[cfg(feature = "hydrate")]
@@ -12,25 +31,7 @@ use std::collections::VecDeque;
 #[cfg(feature = "hydrate")]
 use uuid::Uuid;
 #[cfg(feature = "hydrate")]
-use js_sys::Date;
-#[cfg(feature = "hydrate")]
 use web_sys::{HtmlInputElement, KeyboardEvent, SubmitEvent};
-#[cfg(feature = "hydrate")]
-use futures_util::StreamExt;
-#[cfg(feature = "hydrate")]
-use async_openai_wasm::{
-    types::{
-        ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestSystemMessageArgs,
-        ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs, Model as OpenAIModel,
-    },
-    Client,
-};
-#[cfg(feature = "hydrate")]
-use async_openai_wasm::config::OpenAIConfig;
-#[cfg(feature = "hydrate")]
-use async_openai_wasm::types::{Role, FinishReason};
-#[cfg(feature = "hydrate")]
-use leptos::task::spawn_local;
 
 #[cfg(feature = "hydrate")]
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -43,11 +44,15 @@ pub struct Message {
 
 #[cfg(feature = "hydrate")]
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct MessageContent(pub either::Either<String, Vec<std::collections::HashMap<String, MessageInnerContent>>>);
+pub struct MessageContent(
+    pub either::Either<String, Vec<std::collections::HashMap<String, MessageInnerContent>>>,
+);
 
 #[cfg(feature = "hydrate")]
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct MessageInnerContent(pub either::Either<String, std::collections::HashMap<String, String>>);
+pub struct MessageInnerContent(
+    pub either::Either<String, std::collections::HashMap<String, String>>,
+);
 
 #[cfg(feature = "hydrate")]
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -62,27 +67,40 @@ const DEFAULT_MODEL: &str = "default";
 
 #[cfg(feature = "hydrate")]
 async fn fetch_available_models() -> Result<Vec<OpenAIModel>, String> {
-    leptos::logging::log!("[DEBUG_LOG] fetch_available_models: Starting model fetch from http://localhost:8080/v1");
-    
+    leptos::logging::log!(
+        "[DEBUG_LOG] fetch_available_models: Starting model fetch from http://localhost:8080/v1"
+    );
+
     let config = OpenAIConfig::new().with_api_base("http://localhost:8080/v1".to_string());
     let client = Client::with_config(config);
-    
+
     match client.models().list().await {
         Ok(response) => {
             let model_count = response.data.len();
-            leptos::logging::log!("[DEBUG_LOG] fetch_available_models: Successfully fetched {} models", model_count);
-            
+            leptos::logging::log!(
+                "[DEBUG_LOG] fetch_available_models: Successfully fetched {} models",
+                model_count
+            );
+
             if model_count > 0 {
                 let model_names: Vec<String> = response.data.iter().map(|m| m.id.clone()).collect();
-                leptos::logging::log!("[DEBUG_LOG] fetch_available_models: Available models: {:?}", model_names);
+                leptos::logging::log!(
+                    "[DEBUG_LOG] fetch_available_models: Available models: {:?}",
+                    model_names
+                );
             } else {
-                leptos::logging::log!("[DEBUG_LOG] fetch_available_models: No models returned by server");
+                leptos::logging::log!(
+                    "[DEBUG_LOG] fetch_available_models: No models returned by server"
+                );
             }
-            
+
             Ok(response.data)
-        },
+        }
         Err(e) => {
-            leptos::logging::log!("[DEBUG_LOG] fetch_available_models: Failed to fetch models: {:?}", e);
+            leptos::logging::log!(
+                "[DEBUG_LOG] fetch_available_models: Failed to fetch models: {:?}",
+                e
+            );
             Err(format!("Failed to fetch models: {}", e))
         }
     }
@@ -150,7 +168,7 @@ fn ChatInterface() -> impl IntoView {
     {
         ChatInterfaceImpl()
     }
-    
+
     #[cfg(not(feature = "hydrate"))]
     {
         view! {
@@ -252,7 +270,7 @@ fn ChatInterfaceImpl() -> impl IntoView {
 
             let current_model = selected_model.get_untracked();
             let total_messages = chat_messages.len();
-            
+
             leptos::logging::log!("[DEBUG_LOG] send_message: Preparing request - model: '{}', history_count: {}, total_messages: {}", 
                           current_model, history_count, total_messages);
 
@@ -267,17 +285,17 @@ fn ChatInterfaceImpl() -> impl IntoView {
             // Send request
             let config = OpenAIConfig::new().with_api_base("http://localhost:8080/v1".to_string());
             let client = Client::with_config(config);
-            
+
             leptos::logging::log!("[DEBUG_LOG] send_message: Sending request to http://localhost:8080/v1 with model: '{}'", current_model);
 
             match client.chat().create_stream(request).await {
                 Ok(mut stream) => {
                     leptos::logging::log!("[DEBUG_LOG] send_message: Successfully created stream");
-                    
+
                     let mut assistant_created = false;
                     let mut content_appended = false;
                     let mut chunks_received = 0;
-                    
+
                     while let Some(next) = stream.next().await {
                         match next {
                             Ok(chunk) => {
@@ -335,7 +353,11 @@ fn ChatInterfaceImpl() -> impl IntoView {
                                 }
                             }
                             Err(e) => {
-                                leptos::logging::log!("[DEBUG_LOG] send_message: Stream error after {} chunks: {:?}", chunks_received, e);
+                                leptos::logging::log!(
+                                    "[DEBUG_LOG] send_message: Stream error after {} chunks: {:?}",
+                                    chunks_received,
+                                    e
+                                );
                                 set_messages.update(|msgs| {
                                     msgs.push_back(Message {
                                         id: Uuid::new_v4().to_string(),
@@ -364,7 +386,10 @@ fn ChatInterfaceImpl() -> impl IntoView {
                     leptos::logging::log!("[DEBUG_LOG] send_message: Stream completed successfully, received {} chunks", chunks_received);
                 }
                 Err(e) => {
-                    leptos::logging::log!("[DEBUG_LOG] send_message: Request failed with error: {:?}", e);
+                    leptos::logging::log!(
+                        "[DEBUG_LOG] send_message: Request failed with error: {:?}",
+                        e
+                    );
                     let error_message = Message {
                         id: Uuid::new_v4().to_string(),
                         role: "system".to_string(),
@@ -404,7 +429,8 @@ fn ChatInterfaceImpl() -> impl IntoView {
     };
 
     let messages_list = move || {
-        messages.get()
+        messages
+            .get()
             .into_iter()
             .map(|message| {
                 let role_class = match message.role.as_str() {
@@ -439,7 +465,7 @@ fn ChatInterfaceImpl() -> impl IntoView {
             <h1>"Chat Interface"</h1>
             <div class="model-selector">
                 <label for="model-select">"Model: "</label>
-                <select 
+                <select
                     id="model-select"
                     on:change=on_model_change
                     prop:value=selected_model

crates/leptos-app/src/lib.rs

@@ -10,10 +10,10 @@ pub fn hydrate() {
 
 #[cfg(feature = "ssr")]
 pub fn create_leptos_router() -> axum::Router {
+    use crate::app::*;
     use axum::Router;
     use leptos::prelude::*;
     use leptos_axum::{generate_route_list, LeptosRoutes};
-    use crate::app::*;
 
     let conf = get_configuration(None).unwrap();
     let leptos_options = conf.leptos_options;

crates/leptos-app/src/main.rs

@@ -1,12 +1,11 @@
-
 #[cfg(feature = "ssr")]
 #[tokio::main]
 async fn main() {
     use axum::Router;
     use leptos::logging::log;
     use leptos::prelude::*;
-    use leptos_axum::{generate_route_list, LeptosRoutes};
     use leptos_app::app::*;
+    use leptos_axum::{generate_route_list, LeptosRoutes};
 
     let conf = get_configuration(None).unwrap();
     let addr = conf.leptos_options.site_addr;

crates/llama-runner/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "llama-runner"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+candle-core = { git = "https://github.com/huggingface/candle.git" }
+candle-nn = { git = "https://github.com/huggingface/candle.git" }  
+candle-transformers = { git = "https://github.com/huggingface/candle.git" }
+hf-hub = "0.3"
+tokenizers = "0.20"
+anyhow = "1.0"
+clap = { version = "4.0", features = ["derive", "string"] }
+serde_json = "1.0"
+
+[target.'cfg(target_os = "macos")'.dependencies]
+candle-core = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+candle-nn = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+candle-transformers = { git = "https://github.com/huggingface/candle.git", features = ["metal"] }
+
+[features]
+default = []
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+metal = ["candle-core/metal", "candle-nn/metal", "candle-transformers/metal"]

crates/llama-runner/README.md

@@ -0,0 +1,188 @@
+# Llama Runner
+
+A fast Rust implementation for running Llama and other language models using the Candle deep learning framework. Built on the official Candle examples with optimizations for speed and usability.
+
+## Features
+
+- 🚀 **High Performance**: Metal GPU acceleration on macOS, CUDA support on Linux/Windows
+- 🤖 **Multiple Models**: Supports Llama 3.2, SmolLM2, TinyLlama, and more
+- ⚡ **Fast Inference**: Optimized with F16 precision and KV caching
+- 🎯 **Advanced Sampling**: Top-k, top-p, temperature, and repeat penalty controls  
+- 📊 **Performance Metrics**: Real-time tokens/second reporting
+- 🔧 **Easy CLI**: Simple command-line interface with sensible defaults
+
+## Supported Models
+
+| Model | Size | Command | Description |
+|-------|------|---------|-------------|
+| SmolLM2-135M | 135M | `smollm2-135m` | Tiny, fast model for testing |
+| SmolLM2-360M | 360M | `smollm2-360m` | Small, efficient model |
+| SmolLM2-1.7B | 1.7B | `smollm2-1.7b` | Balanced performance/speed |
+| Llama-3.2-1B | 1B | `llama-3.2-1b` | Meta's compact model |
+| Llama-3.2-3B | 3B | `llama-3.2-3b` | Larger Llama model |
+| TinyLlama-1.1B | 1.1B | `tinyllama-1.1b-chat` | Chat-optimized small model |
+
+Add `-instruct` suffix for instruction-tuned variants (e.g., `smollm2-135m-instruct`).
+
+## Installation
+
+```bash
+# Clone the repository
+git clone <repository-url>
+cd llama-runner
+
+# Build with GPU acceleration (recommended)
+cargo build --release --features metal  # macOS
+cargo build --release --features cuda   # Linux/Windows with NVIDIA GPU
+
+# CPU-only build
+cargo build --release
+```
+
+## Quick Start
+
+```bash
+# Fast inference with GPU acceleration
+cargo run --features metal -- --prompt "What is quantum computing?"
+
+# Specify a model and parameters
+cargo run --features metal -- \
+  --prompt "Write a short story about space exploration" \
+  --model smollm2-360m \
+  --max-tokens 100 \
+  --temperature 0.8
+
+# Use CPU (slower but works everywhere)
+cargo run -- --prompt "Hello, world!" --model smollm2-135m --cpu
+```
+
+## Usage Examples
+
+### Basic Text Generation
+```bash
+# Simple completion
+cargo run --features metal -- --prompt "The capital of France is"
+
+# Creative writing with higher temperature
+cargo run --features metal -- \
+  --prompt "Once upon a time" \
+  --temperature 1.0 \
+  --max-tokens 200
+```
+
+### Advanced Sampling
+```bash
+# Top-k and top-p sampling
+cargo run --features metal -- \
+  --prompt "Explain artificial intelligence" \
+  --top-k 40 \
+  --top-p 0.9 \
+  --temperature 0.7
+
+# Reduce repetition
+cargo run --features metal -- \
+  --prompt "List the benefits of renewable energy" \
+  --repeat-penalty 1.2 \
+  --repeat-last-n 64
+```
+
+### Different Models
+```bash
+# Ultra-fast with tiny model
+cargo run --features metal -- \
+  --prompt "Quick test" \
+  --model smollm2-135m
+
+# Better quality with larger model
+cargo run --features metal -- \
+  --prompt "Explain quantum physics" \
+  --model llama-3.2-1b \
+  --max-tokens 150
+```
+
+## Command-Line Options
+
+| Option | Short | Default | Description |
+|--------|-------|---------|-------------|
+| `--prompt` | `-p` | "The capital of France is" | Input prompt |
+| `--model` | `-m` | `smollm2-135m` | Model to use |
+| `--max-tokens` | `-n` | 100 | Maximum tokens to generate |
+| `--temperature` | `-t` | 0.8 | Sampling temperature (0.0 = deterministic) |
+| `--top-k` | | None | Top-k sampling |
+| `--top-p` | | None | Top-p (nucleus) sampling |
+| `--seed` | | 299792458 | Random seed for reproducibility |
+| `--repeat-penalty` | | 1.1 | Repetition penalty (1.0 = no penalty) |
+| `--repeat-last-n` | | 128 | Context window for repeat penalty |
+| `--cpu` | | false | Force CPU usage |
+| `--dtype` | | f16 | Data type: f16, bf16, f32 |
+| `--no-kv-cache` | | false | Disable key-value caching |
+
+## Performance
+
+Typical performance on Apple M2 with Metal acceleration:
+
+| Model | Size | Speed | Memory |
+|-------|------|-------|--------|
+| SmolLM2-135M | 135M | ~100 tok/s | ~500MB |
+| SmolLM2-360M | 360M | ~80 tok/s | ~1GB |
+| SmolLM2-1.7B | 1.7B | ~50 tok/s | ~3GB |
+| Llama-3.2-1B | 1B | ~40 tok/s | ~2GB |
+
+## Requirements
+
+- **Rust**: 1.70+ (latest stable recommended)
+- **Memory**: 2-8GB RAM depending on model size
+- **Storage**: 1-10GB for model weights
+- **Network**: Internet connection for first-time model download
+- **GPU** (optional): Metal on macOS, CUDA on Linux/Windows
+
+## GPU Support
+
+### macOS (Metal)
+```bash
+cargo run --features metal -- [options]
+```
+
+### Linux/Windows (CUDA)
+```bash
+cargo run --features cuda -- [options]  
+```
+
+### CPU Only
+```bash
+cargo run -- --cpu [options]
+```
+
+## Model Downloads
+
+Models are automatically downloaded from HuggingFace Hub on first use and cached locally. Download times:
+
+- SmolLM2-135M: ~1 minute
+- SmolLM2-360M: ~2 minutes  
+- Llama-3.2-1B: ~5 minutes
+- Larger models: 10+ minutes
+
+## Troubleshooting
+
+### Slow Performance
+- Use `--features metal` on macOS or `--features cuda` on Linux/Windows
+- Try smaller models like `smollm2-135m` for faster inference
+- Ensure sufficient RAM for your chosen model
+
+### Out of Memory
+- Use `--cpu` to use system RAM instead of GPU memory
+- Try smaller models or reduce `--max-tokens`
+- Use `--dtype f32` if f16 causes issues
+
+### Model Download Issues
+- Check internet connection
+- Some models may require HuggingFace Hub authentication
+- Verify sufficient disk space in `~/.cache/huggingface/`
+
+## Contributing
+
+Contributions welcome! This project is based on the [Candle](https://github.com/huggingface/candle) framework by HuggingFace.
+
+## License
+
+MIT License - see LICENSE file for details.

crates/llama-runner/src/lib.rs

@@ -0,0 +1,7 @@
+pub mod llama_api;
+
+use clap::ValueEnum;
+pub use llama_api::{run_llama_inference, LlamaInferenceConfig, WhichModel};
+
+// Re-export constants and types that might be needed
+pub const EOS_TOKEN: &str = "</s>";

crates/llama-runner/src/llama_api.rs

@@ -0,0 +1,333 @@
+use crate::EOS_TOKEN;
+use anyhow::{bail, Error as E};
+use candle_core::{utils, DType, Device, Tensor};
+use candle_nn::VarBuilder;
+use candle_transformers::generation::{LogitsProcessor, Sampling};
+use candle_transformers::models::llama as model;
+use candle_transformers::models::llama::{Llama, LlamaConfig};
+use clap::ValueEnum;
+use hf_hub::api::sync::Api;
+use hf_hub::{Repo, RepoType};
+use std::sync::mpsc::{self, Receiver};
+
+#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum, Default)]
+pub enum WhichModel {
+    #[value(name = "llama-3.2-1b")]
+    #[default]
+    Llama32_1B,
+    #[value(name = "llama-3.2-1b-instruct")]
+    Llama32_1BInstruct,
+    #[value(name = "llama-3.2-3b")]
+    Llama32_3B,
+    #[value(name = "llama-3.2-3b-instruct")]
+    Llama32_3BInstruct,
+    #[value(name = "smollm2-135m")]
+    SmolLM2_135M,
+    #[value(name = "smollm2-135m-instruct")]
+    SmolLM2_135MInstruct,
+    #[value(name = "smollm2-360m")]
+    SmolLM2_360M,
+    #[value(name = "smollm2-360m-instruct")]
+    SmolLM2_360MInstruct,
+    #[value(name = "smollm2-1.7b")]
+    SmolLM2_1_7B,
+    #[value(name = "smollm2-1.7b-instruct")]
+    SmolLM2_1_7BInstruct,
+    #[value(name = "tinyllama-1.1b-chat")]
+    TinyLlama1_1BChat,
+}
+
+#[derive(Debug, Clone)]
+pub struct LlamaInferenceConfig {
+    pub prompt: String,
+
+    pub model: WhichModel,
+    pub cpu: bool,
+    pub temperature: f64,
+    pub top_p: Option<f64>,
+    pub top_k: Option<usize>,
+    pub seed: u64,
+    pub max_tokens: usize,
+    pub no_kv_cache: bool,
+    pub dtype: Option<String>,
+    pub model_id: Option<String>,
+    pub revision: Option<String>,
+    pub use_flash_attn: bool,
+    pub repeat_penalty: f32,
+    pub repeat_last_n: usize,
+}
+
+impl Default for LlamaInferenceConfig {
+    fn default() -> Self {
+        Self {
+            // Leave prompt empty by default; let call sites set it.
+            prompt: String::new(),
+
+            // Keep your existing model choice; swap at call-site if needed.
+            model: WhichModel::Llama32_1BInstruct,
+
+            // Prefer GPU if available.
+            cpu: false,
+
+            // Sampling: balanced + stable
+            temperature: 0.7,
+            top_p: Some(0.95),
+            top_k: Some(50),
+
+            // Reproducible by default; override for variability.
+            seed: 42,
+
+            // Don’t run unbounded generations.
+            max_tokens: 512,
+
+            // Performance flags
+            no_kv_cache: false,   // keep cache ON for speed
+            use_flash_attn: true, // great speed boost if supported
+
+            // Precision: bf16 is a good default on Ampere+; fallback to fp16 if needed.
+            dtype: Some("bf16".to_string()),
+
+            // Optional model source pinning (None = app defaults)
+            model_id: None,
+            revision: None,
+
+            // Anti-repeat heuristics
+            repeat_penalty: 1.15,
+            repeat_last_n: 128,
+        }
+    }
+}
+
+fn device(cpu: bool) -> anyhow::Result<Device> {
+    if cpu {
+        Ok(Device::Cpu)
+    } else if utils::cuda_is_available() {
+        Ok(Device::new_cuda(0)?)
+    } else if utils::metal_is_available() {
+        Ok(Device::new_metal(0)?)
+    } else {
+        Ok(Device::Cpu)
+    }
+}
+
+fn hub_load_safetensors(
+    api: &hf_hub::api::sync::ApiRepo,
+    json_file: &str,
+) -> anyhow::Result<Vec<std::path::PathBuf>> {
+    let json_file = api.get(json_file)?;
+    let json_file = std::fs::File::open(json_file)?;
+    let json: serde_json::Value = serde_json::from_reader(&json_file)?;
+    let weight_map = match json.get("weight_map") {
+        None => bail!("no weight map in {json_file:?}"),
+        Some(serde_json::Value::Object(map)) => map,
+        Some(_) => bail!("weight map in {json_file:?} is not a map"),
+    };
+    let mut safetensors_files = std::collections::HashSet::new();
+    for value in weight_map.values() {
+        if let Some(file) = value.as_str() {
+            safetensors_files.insert(file.to_string());
+        }
+    }
+    let safetensors_files = safetensors_files
+        .iter()
+        .map(|v| api.get(v))
+        .collect::<anyhow::Result<Vec<_>, _>>()?;
+    Ok(safetensors_files)
+}
+
+pub fn run_llama_inference(
+    cfg: LlamaInferenceConfig,
+) -> anyhow::Result<Receiver<anyhow::Result<String>>, anyhow::Error> {
+    // ---- Device & dtype -----------------------------------------------------
+    let device = device(cfg.cpu)?;
+    println!("Device: {:?}", device);
+
+    let dtype = match cfg.dtype.as_deref() {
+        Some("f16") => DType::F16,
+        Some("bf16") => DType::BF16,
+        Some("f32") => DType::F32,
+        Some(dtype) => bail!("Unsupported dtype {dtype}"),
+        None => DType::F16,
+    };
+    println!("Using dtype: {:?}", dtype);
+
+    // ---- Load model & tokenizer --------------------------------------------
+    let (llama, tokenizer, mut cache) = {
+        let api = Api::new()?;
+        let model_id = cfg.model_id.clone().unwrap_or_else(|| {
+            match cfg.model {
+                WhichModel::Llama32_1B => "meta-llama/Llama-3.2-1B",
+                WhichModel::Llama32_1BInstruct => "meta-llama/Llama-3.2-1B-Instruct",
+                WhichModel::Llama32_3B => "meta-llama/Llama-3.2-3B",
+                WhichModel::Llama32_3BInstruct => "meta-llama/Llama-3.2-3B-Instruct",
+                WhichModel::SmolLM2_135M => "HuggingFaceTB/SmolLM2-135M",
+                WhichModel::SmolLM2_135MInstruct => "HuggingFaceTB/SmolLM2-135M-Instruct",
+                WhichModel::SmolLM2_360M => "HuggingFaceTB/SmolLM2-360M",
+                WhichModel::SmolLM2_360MInstruct => "HuggingFaceTB/SmolLM2-360M-Instruct",
+                WhichModel::SmolLM2_1_7B => "HuggingFaceTB/SmolLM2-1.7B",
+                WhichModel::SmolLM2_1_7BInstruct => "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+                WhichModel::TinyLlama1_1BChat => "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+            }
+            .to_string()
+        });
+        println!("Loading model: {}", model_id);
+        let revision = cfg.revision.clone().unwrap_or("main".to_string());
+        let api = api.repo(Repo::with_revision(model_id, RepoType::Model, revision));
+
+        let tokenizer_filename = api.get("tokenizer.json")?;
+        let config_filename = api.get("config.json")?;
+        let config: LlamaConfig = serde_json::from_slice(&std::fs::read(config_filename)?)?;
+        let config = config.into_config(cfg.use_flash_attn);
+
+        let filenames = match cfg.model {
+            WhichModel::Llama32_3B | WhichModel::Llama32_3BInstruct => {
+                hub_load_safetensors(&api, "model.safetensors.index.json")?
+            }
+            _ => vec![api.get("model.safetensors")?],
+        };
+
+        let cache = model::Cache::new(!cfg.no_kv_cache, dtype, &config, &device)?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, dtype, &device)? };
+        let llama = Llama::load(vb, &config)?;
+        let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+        (llama, tokenizer, cache)
+    };
+
+    // ---- Prepare prompt & sampler ------------------------------------------
+    let eos_token_id = tokenizer
+        .token_to_id(EOS_TOKEN)
+        .map(model::LlamaEosToks::Single);
+
+    let mut tokens = tokenizer
+        .encode(cfg.prompt.as_str(), true)
+        .map_err(E::msg)?
+        .get_ids()
+        .to_vec();
+
+    println!("Starting inference...");
+
+    let mut logits_processor = {
+        let temperature = cfg.temperature;
+        let sampling = if temperature <= 0. {
+            Sampling::ArgMax
+        } else {
+            match (cfg.top_k, cfg.top_p) {
+                (None, None) => Sampling::All { temperature },
+                (Some(k), None) => Sampling::TopK { k, temperature },
+                (None, Some(p)) => Sampling::TopP { p, temperature },
+                (Some(k), Some(p)) => Sampling::TopKThenTopP { k, p, temperature },
+            }
+        };
+        LogitsProcessor::from_sampling(cfg.seed, sampling)
+    };
+
+    // Channel for streaming decoded fragments to the caller.
+    let (tx, rx) = mpsc::channel::<anyhow::Result<String>>();
+
+    // ---- Spawn generation thread -------------------------------------------
+    std::thread::spawn(move || {
+        let start_gen = std::time::Instant::now();
+        let mut index_pos = 0usize;
+        let mut token_generated = 0usize;
+
+        for index in 0..cfg.max_tokens {
+            // Use KV-cache for single-token step after the first pass.
+            let (context_size, context_index) = if cache.use_kv_cache && index > 0 {
+                (1, index_pos)
+            } else {
+                (tokens.len(), 0)
+            };
+
+            let ctxt = &tokens[tokens.len().saturating_sub(context_size)..];
+            let input = match Tensor::new(ctxt, &device).and_then(|t| t.unsqueeze(0)) {
+                Ok(t) => t,
+                Err(e) => {
+                    let _ = tx.send(Err(e.into()));
+                    break;
+                }
+            };
+
+            let logits = match llama.forward(&input, context_index, &mut cache) {
+                Ok(l) => l,
+                Err(e) => {
+                    let _ = tx.send(Err(e.into()));
+                    break;
+                }
+            };
+            let logits = match logits.squeeze(0) {
+                Ok(l) => l,
+                Err(e) => {
+                    let _ = tx.send(Err(e.into()));
+                    break;
+                }
+            };
+
+            let logits = if cfg.repeat_penalty == 1. {
+                logits
+            } else {
+                let start_at = tokens.len().saturating_sub(cfg.repeat_last_n);
+                match candle_transformers::utils::apply_repeat_penalty(
+                    &logits,
+                    cfg.repeat_penalty,
+                    &tokens[start_at..],
+                ) {
+                    Ok(l) => l,
+                    Err(e) => {
+                        let _ = tx.send(Err(e.into()));
+                        break;
+                    }
+                }
+            };
+
+            index_pos += ctxt.len();
+
+            let next_token = match logits_processor.sample(&logits) {
+                Ok(t) => t,
+                Err(e) => {
+                    let _ = tx.send(Err(e.into()));
+                    break;
+                }
+            };
+
+            token_generated += 1;
+            tokens.push(next_token);
+
+            // Early stop on EOS.
+            let stop = match eos_token_id {
+                Some(model::LlamaEosToks::Single(eos_tok_id)) => next_token == eos_tok_id,
+                Some(model::LlamaEosToks::Multiple(ref eos_ids)) => eos_ids.contains(&next_token),
+                None => false,
+            };
+            if stop {
+                break;
+            }
+
+            // Decode this token's text and stream it out.
+            match tokenizer.decode(&[next_token], false) {
+                Ok(text) => {
+                    if !text.is_empty() {
+                        // Best-effort send; if receiver is gone, just stop.
+                        if tx.send(Ok(text)).is_err() {
+                            break;
+                        }
+                    }
+                }
+                Err(e) => {
+                    let _ = tx.send(Err(anyhow::anyhow!("{}", e)));
+                    break;
+                }
+            }
+        }
+
+        // Optional: final stats as a debug line (not sent through the stream).
+        let dt = start_gen.elapsed();
+        eprintln!(
+            "[llama-runner] {} tokens generated ({:.2} tokens/s)",
+            token_generated,
+            token_generated as f64 / dt.as_secs_f64(),
+        );
+        // Dropping tx closes the stream.
+    });
+
+    Ok(rx)
+}

crates/llama-runner/src/llama_cli.rs

@@ -0,0 +1,108 @@
+use crate::llama_api::{run_llama_inference, LlamaInferenceConfig, WhichModel};
+use clap::Parser;
+use std::io::Write;
+
+#[derive(Parser, Debug, Default)]
+#[command(author, version, about = "Fast Llama inference with Candle", long_about = None)]
+struct Args {
+    /// The prompt to generate text from
+    #[arg(short, long, default_value = "The capital of France is")]
+    prompt: String,
+
+    /// The model to use
+    #[arg(short, long, default_value = "llama-3.2-1b-instruct")]
+    model: WhichModel,
+
+    /// Run on CPU rather than GPU
+    #[arg(long)]
+    cpu: bool,
+
+    /// The temperature used to generate samples
+    #[arg(short, long, default_value_t = 0.8)]
+    temperature: f64,
+
+    /// Nucleus sampling probability cutoff
+    #[arg(long)]
+    top_p: Option<f64>,
+
+    /// Only sample among the top K samples
+    #[arg(long)]
+    top_k: Option<usize>,
+
+    /// The seed to use when generating random samples
+    #[arg(long, default_value_t = 299792458)]
+    seed: u64,
+
+    /// The length of the sample to generate (in tokens)
+    #[arg(short = 'n', long, default_value_t = 100)]
+    max_tokens: usize,
+
+    /// Disable the key-value cache
+    #[arg(long)]
+    no_kv_cache: bool,
+
+    /// Use different dtype than f16
+    #[arg(long)]
+    dtype: Option<String>,
+
+    /// Custom model ID from HuggingFace Hub
+    #[arg(long)]
+    model_id: Option<String>,
+
+    /// Model revision
+    #[arg(long)]
+    revision: Option<String>,
+
+    /// Use flash attention
+    #[arg(long)]
+    use_flash_attn: bool,
+
+    /// Penalty to be applied for repeating tokens, 1. means no penalty
+    #[arg(long, default_value_t = 1.1)]
+    repeat_penalty: f32,
+
+    /// The context size to consider for the repeat penalty
+    #[arg(long, default_value_t = 128)]
+    repeat_last_n: usize,
+}
+
+impl Into<LlamaInferenceConfig> for Args {
+    fn into(self) -> LlamaInferenceConfig {
+        LlamaInferenceConfig {
+            prompt: self.prompt,
+            model: self.model,
+            cpu: self.cpu,
+            temperature: self.temperature,
+            top_p: self.top_p,
+            top_k: self.top_k,
+            seed: self.seed,
+            max_tokens: self.max_tokens,
+            no_kv_cache: self.no_kv_cache,
+            dtype: self.dtype,
+            model_id: self.model_id,
+            revision: self.revision,
+            use_flash_attn: self.use_flash_attn,
+            repeat_penalty: self.repeat_penalty,
+            repeat_last_n: self.repeat_last_n,
+        }
+    }
+}
+
+pub fn run_cli() -> anyhow::Result<()> {
+    let args = Args::parse();
+    let cfg = args.into();
+    let rx = run_llama_inference(cfg)?;
+    for msg in rx {
+        match msg {
+            Ok(tok) => {
+                print!("{tok}");
+                let _ = std::io::stdout().flush(); // <- force it out now
+            }
+            Err(e) => {
+                eprintln!("generation error: {e}");
+                break;
+            }
+        }
+    }
+    Ok(())
+}

crates/llama-runner/src/main.rs

@@ -0,0 +1,19 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+mod llama_api;
+mod llama_cli;
+
+use anyhow::Result;
+use clap::{Parser, ValueEnum};
+
+use std::io::Write;
+
+use crate::llama_cli::run_cli;
+
+const EOS_TOKEN: &str = "</s>";
+
+fn main() -> Result<()> {
+    run_cli()
+}

crates/predict-otron-9000/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "predict-otron-9000"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2024"
 
 [[bin]]

crates/predict-otron-9000/src/config.rs

@@ -1,7 +1,9 @@
 use serde::{Deserialize, Serialize};
 use std::env;
+use tracing::info;
+use tracing::log::error;
 
-#[derive(Debug, Clone, Deserialize, Serialize)]
+#[derive(Serialize, Deserialize, Clone, Debug)]
 #[serde(rename_all = "camelCase")]
 pub struct ServerConfig {
     #[serde(default = "default_server_host")]
@@ -10,14 +12,16 @@ pub struct ServerConfig {
     pub server_port: u16,
     pub server_mode: ServerMode,
     #[serde(default)]
-    pub services: Services,
+    pub services: Option<Services>,
 }
 
 fn default_server_host() -> String {
     "127.0.0.1".to_string()
 }
 
-fn default_server_port() -> u16 { 8080 }
+fn default_server_port() -> u16 {
+    8080
+}
 
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 #[serde(rename_all = "PascalCase")]
@@ -34,17 +38,15 @@ impl Default for ServerMode {
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Services {
-    #[serde(default = "inference_service_url")]
-    pub inference_url: String,
-    #[serde(default = "embeddings_service_url")]
-    pub embeddings_url: String,
+    pub inference_url: Option<String>,
+    pub embeddings_url: Option<String>,
 }
 
 impl Default for Services {
     fn default() -> Self {
         Self {
-            inference_url: inference_service_url(),
-            embeddings_url: embeddings_service_url(),
+            inference_url: None,
+            embeddings_url: None,
         }
     }
 }
@@ -63,7 +65,7 @@ impl Default for ServerConfig {
             server_host: "127.0.0.1".to_string(),
             server_port: 8080,
             server_mode: ServerMode::Standalone,
-            services: Services::default(),
+            services: Some(Services::default()),
         }
     }
 }
@@ -73,21 +75,19 @@ impl ServerConfig {
     /// Falls back to default (Local mode) if not set or invalid
     pub fn from_env() -> Self {
         match env::var("SERVER_CONFIG") {
-            Ok(config_str) => {
-                match serde_json::from_str::<ServerConfig>(&config_str) {
-                    Ok(config) => {
-                        tracing::info!("Loaded server configuration: {:?}", config);
-                        config
-                    }
-                    Err(e) => {
-                        tracing::warn!(
-                            "Failed to parse SERVER_CONFIG environment variable: {}. Using default configuration.", 
-                            e
-                        );
-                        ServerConfig::default()
-                    }
+            Ok(config_str) => match serde_json::from_str::<ServerConfig>(&config_str) {
+                Ok(config) => {
+                    tracing::info!("Loaded server configuration: {:?}", config);
+                    config
                 }
-            }
+                Err(e) => {
+                    tracing::warn!(
+                        "Failed to parse SERVER_CONFIG environment variable: {}. Using default configuration.",
+                        e
+                    );
+                    ServerConfig::default()
+                }
+            },
             Err(_) => {
                 tracing::info!("SERVER_CONFIG not set, Standalone mode active");
                 ServerConfig::default()
@@ -96,18 +96,52 @@ impl ServerConfig {
     }
 
     /// Check if the server should run in high availability mode
-    pub fn is_high_availability(&self) -> bool {
-        self.server_mode == ServerMode::HighAvailability
+    pub fn is_high_availability(&self) -> Result<bool, std::io::Error> {
+        if self.server_mode == ServerMode::HighAvailability {
+            let services_well_defined: bool = self.clone().services.is_some();
+
+            let inference_url_well_defined: bool =
+                services_well_defined && self.clone().services.unwrap().inference_url.is_some();
+
+            let embeddings_well_defined: bool =
+                services_well_defined && self.clone().services.unwrap().embeddings_url.is_some();
+
+            let is_well_defined_for_ha =
+                services_well_defined && inference_url_well_defined && embeddings_well_defined;
+
+            if !is_well_defined_for_ha {
+                let config_string = serde_json::to_string_pretty(&self).unwrap();
+                error!(
+                    "HighAvailability mode configured but services not well defined! \n## Config Used:\n {}",
+                    config_string
+                );
+                let err = std::io::Error::new(
+                    std::io::ErrorKind::Other,
+                    "HighAvailability mode configured but services not well defined!",
+                );
+                return Err(err);
+            }
+        }
+
+        Ok(self.server_mode == ServerMode::HighAvailability)
     }
 
     /// Get the inference service URL for proxying
-    pub fn inference_url(&self) -> &str {
-        &self.services.inference_url
+    pub fn inference_url(&self) -> Option<String> {
+        if self.services.is_some() {
+            self.services.clone()?.inference_url
+        } else {
+            None
+        }
     }
 
     /// Get the embeddings service URL for proxying
-    pub fn embeddings_url(&self) -> &str {
-        &self.services.embeddings_url
+    pub fn embeddings_url(&self) -> Option<String> {
+        if self.services.is_some() {
+            self.services.clone()?.embeddings_url
+        } else {
+            None
+        }
     }
 }
 
@@ -119,7 +153,7 @@ mod tests {
     fn test_default_config() {
         let config = ServerConfig::default();
         assert_eq!(config.server_mode, ServerMode::Standalone);
-        assert!(!config.is_high_availability());
+        assert!(!config.is_high_availability().unwrap());
     }
 
     #[test]
@@ -134,23 +168,26 @@ mod tests {
 
         let config: ServerConfig = serde_json::from_str(config_json).unwrap();
         assert_eq!(config.server_mode, ServerMode::HighAvailability);
-        assert!(config.is_high_availability());
-        assert_eq!(config.inference_url(), "http://inference-service:8080");
-        assert_eq!(config.embeddings_url(), "http://embeddings-service:8080");
+        assert!(config.is_high_availability().unwrap());
+        assert_eq!(
+            config.inference_url().unwrap(),
+            "http://inference-service:8080"
+        );
+        assert_eq!(
+            config.embeddings_url().unwrap(),
+            "http://embeddings-service:8080"
+        );
     }
 
     #[test]
     fn test_local_mode_config() {
         let config_json = r#"{
-            "serverMode": "Local"
+            "serverMode": "Standalone"
         }"#;
 
         let config: ServerConfig = serde_json::from_str(config_json).unwrap();
         assert_eq!(config.server_mode, ServerMode::Standalone);
-        assert!(!config.is_high_availability());
-        // Should use default URLs
-        assert_eq!(config.inference_url(), "http://inference-service:8080");
-        assert_eq!(config.embeddings_url(), "http://embeddings-service:8080");
+        assert!(!config.is_high_availability().unwrap());
     }
 
     #[test]
@@ -164,17 +201,26 @@ mod tests {
         }"#;
 
         let config: ServerConfig = serde_json::from_str(config_json).unwrap();
-        assert_eq!(config.inference_url(), "http://custom-inference:9000");
-        assert_eq!(config.embeddings_url(), "http://custom-embeddings:9001");
+        assert_eq!(
+            config.inference_url().unwrap(),
+            "http://custom-inference:9000"
+        );
+        assert_eq!(
+            config.embeddings_url().unwrap(),
+            "http://custom-embeddings:9001"
+        );
     }
 
     #[test]
-    fn test_minimal_high_availability_config() {
+    fn test_minimal_high_availability_config_error() {
         let config_json = r#"{"serverMode": "HighAvailability"}"#;
         let config: ServerConfig = serde_json::from_str(config_json).unwrap();
-        assert!(config.is_high_availability());
-        // Should use default URLs
-        assert_eq!(config.inference_url(), "http://inference-service:8080");
-        assert_eq!(config.embeddings_url(), "http://embeddings-service:8080");
+
+        let is_high_availability = config.is_high_availability();
+
+        assert!(is_high_availability.is_err());
+        // // Should use default URLs
+        // assert_eq!(config.inference_url().unwrap(), "http://inference-service:8080");
+        // assert_eq!(config.embeddings_url().unwrap(), "http://embeddings-service:8080");
     }
-}
+}

crates/predict-otron-9000/src/main.rs

@@ -1,7 +1,9 @@
 mod config;
 mod middleware;
 mod proxy;
+mod standalone;
 
+use crate::standalone::create_standalone_router;
 use axum::response::IntoResponse;
 use axum::routing::get;
 use axum::{Router, http::Uri, response::Html, serve};
@@ -11,6 +13,7 @@ use middleware::{MetricsLayer, MetricsLoggerFuture, MetricsStore};
 use proxy::create_proxy_router;
 use rust_embed::Embed;
 use std::env;
+use std::path::Component::ParentDir;
 use tokio::net::TcpListener;
 use tower_http::classify::ServerErrorsFailureClass::StatusCode;
 use tower_http::cors::{Any, CorsLayer};
@@ -49,44 +52,19 @@ async fn main() {
     let default_host = server_config.server_host.clone();
     let default_port = server_config.server_port;
 
-    // Create router based on server mode
-    let service_router = if server_config.clone().is_high_availability() {
-        tracing::info!("Running in HighAvailability mode - proxying to external services");
-        tracing::info!("  Inference service URL: {}", server_config.inference_url());
-        tracing::info!(
-            "  Embeddings service URL: {}",
-            server_config.embeddings_url()
-        );
-
-        // Use proxy router that forwards requests to external services
-        create_proxy_router(server_config.clone())
-    } else {
-        tracing::info!("Running in Standalone mode - using embedded services");
-
-        // Create unified router by merging embeddings and inference routers (existing behavior)
-        let embeddings_router = embeddings_engine::create_embeddings_router();
-
-        // Create AppState with correct model configuration
-        use inference_engine::Which;
-        use inference_engine::server::{PipelineArgs, build_pipeline};
-        let mut pipeline_args = PipelineArgs::default();
-        pipeline_args.model_id = "google/gemma-3-1b-it".to_string();
-        pipeline_args.which = Which::InstructV3_1B;
-
-        let text_generation = build_pipeline(pipeline_args.clone());
-        let app_state = AppState {
-            text_generation: std::sync::Arc::new(tokio::sync::Mutex::new(text_generation)),
-            model_id: "google/gemma-3-1b-it".to_string(),
-            build_args: pipeline_args,
-        };
-
-        // Get the inference router directly from the inference engine
-        let inference_router = inference_engine::create_router(app_state);
-
-        // Merge the local routers
-        Router::new()
-            .merge(embeddings_router)
-            .merge(inference_router)
+    let service_router = match server_config.clone().is_high_availability() {
+        Ok(is_ha) => {
+            if is_ha {
+                log_config(server_config.clone());
+                create_proxy_router(server_config.clone())
+            } else {
+                log_config(server_config.clone());
+                create_standalone_router(server_config)
+            }
+        }
+        Err(error) => {
+            panic!("{}", error);
+        }
     };
 
     // Create CORS layer
@@ -135,5 +113,25 @@ async fn main() {
     serve(listener, app).await.unwrap();
 }
 
+fn log_config(config: ServerConfig) {
+    match config.is_high_availability() {
+        Ok(is_high) => {
+            if is_high {
+                tracing::info!("Running in HighAvailability mode - proxying to external services");
+                tracing::info!("Inference service URL: {}", config.inference_url().unwrap());
+                tracing::info!(
+                    "Embeddings service URL: {}",
+                    config.embeddings_url().unwrap()
+                );
+            } else {
+                tracing::info!("Running in Standalone mode");
+            }
+        }
+        Err(error) => {
+            panic!("{}", error);
+        }
+    }
+}
+
 // Chat completions handler that properly uses the inference server crate's error handling
 // This function is no longer needed as we're using the inference_engine router directly

crates/predict-otron-9000/src/middleware/metrics.rs

@@ -2,6 +2,8 @@ use axum::{
     extract::MatchedPath,
     http::{Request, Response},
 };
+use std::fmt;
+use std::task::ready;
 use std::{
     future::Future,
     pin::Pin,
@@ -12,8 +14,6 @@ use std::{
 use tokio::sync::Mutex;
 use tower::{Layer, Service};
 use tracing::{debug, info};
-use std::task::ready;
-use std::fmt;
 
 /// Performance metrics for a specific endpoint
 #[derive(Debug, Clone, Default)]
@@ -33,16 +33,16 @@ impl EndpointMetrics {
     pub fn add_response_time(&mut self, time_ms: u64) {
         self.count += 1;
         self.total_time_ms += time_ms;
-        
+
         if self.min_time_ms == 0 || time_ms < self.min_time_ms {
             self.min_time_ms = time_ms;
         }
-        
+
         if time_ms > self.max_time_ms {
             self.max_time_ms = time_ms;
         }
     }
-    
+
     /// Get the average response time in milliseconds
     pub fn avg_time_ms(&self) -> f64 {
         if self.count == 0 {
@@ -51,12 +51,15 @@ impl EndpointMetrics {
             self.total_time_ms as f64 / self.count as f64
         }
     }
-    
+
     /// Get a human-readable summary of the metrics
     pub fn summary(&self) -> String {
         format!(
             "requests: {}, avg: {:.2}ms, min: {}ms, max: {}ms",
-            self.count, self.avg_time_ms(), self.min_time_ms, self.max_time_ms
+            self.count,
+            self.avg_time_ms(),
+            self.min_time_ms,
+            self.max_time_ms
         )
     }
 }
@@ -75,14 +78,16 @@ impl MetricsStore {
             endpoints: Arc::new(Mutex::new(std::collections::HashMap::new())),
         }
     }
-    
+
     /// Record a request's timing information
     pub async fn record(&self, path: String, time_ms: u64) {
         let mut endpoints = self.endpoints.lock().await;
-        let metrics = endpoints.entry(path).or_insert_with(EndpointMetrics::default);
+        let metrics = endpoints
+            .entry(path)
+            .or_insert_with(EndpointMetrics::default);
         metrics.add_response_time(time_ms);
     }
-    
+
     /// Get metrics for all endpoints
     pub async fn get_all(&self) -> Vec<(String, EndpointMetrics)> {
         let endpoints = self.endpoints.lock().await;
@@ -91,12 +96,12 @@ impl MetricsStore {
             .map(|(k, v)| (k.clone(), v.clone()))
             .collect()
     }
-    
+
     /// Log a summary of all metrics
     pub async fn log_summary(&self) {
         let metrics = self.get_all().await;
         info!("Performance metrics summary:");
-        
+
         for (path, metric) in metrics {
             info!("  {}: {}", path, metric.summary());
         }
@@ -163,26 +168,28 @@ where
         } else {
             req.uri().path().to_string()
         };
-        
+
         let method = req.method().clone();
         let start = Instant::now();
         let metrics_store = self.metrics_store.clone();
-        
+
         let future = self.inner.call(req);
-        
+
         Box::pin(async move {
             let response = future.await?;
-            
+
             let time = start.elapsed();
             let status = response.status();
             let time_ms = time.as_millis() as u64;
-            
+
             // Record the timing in our metrics store
-            metrics_store.record(format!("{} {}", method, path), time_ms).await;
-            
+            metrics_store
+                .record(format!("{} {}", method, path), time_ms)
+                .await;
+
             // Log the request timing
             debug!("{} {} {} - {} ms", method, path, status, time_ms);
-            
+
             Ok(response)
         })
     }
@@ -214,7 +221,7 @@ impl Future for MetricsLoggerFuture {
                 metrics_store.log_summary().await;
             });
         }
-        
+
         Poll::Pending
     }
-}
+}

crates/predict-otron-9000/src/middleware/mod.rs

@@ -1,7 +1,3 @@
 pub mod metrics;
 
-pub use metrics::{
-    MetricsStore,
-    MetricsLoggerFuture,
-    MetricsLayer,
-};
+pub use metrics::{MetricsLayer, MetricsLoggerFuture, MetricsStore};

crates/predict-otron-9000/src/proxy.rs

@@ -1,10 +1,10 @@
 use axum::{
+    Router,
     body::Body,
     extract::{Request, State},
     http::{HeaderMap, Method, StatusCode, Uri},
     response::{IntoResponse, Response},
     routing::{get, post},
-    Router,
 };
 use reqwest::Client;
 use serde_json::Value;
@@ -47,10 +47,16 @@ async fn proxy_chat_completions(
     headers: HeaderMap,
     body: Body,
 ) -> Result<Response, StatusCode> {
-    let target_url = format!("{}/v1/chat/completions", proxy_client.config.inference_url());
-    
+    let target_url = format!(
+        "{}/v1/chat/completions",
+        proxy_client
+            .config
+            .inference_url()
+            .expect("Invalid Configuration")
+    );
+
     tracing::info!("Proxying chat completions request to: {}", target_url);
-    
+
     // Extract body as bytes
     let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
         Ok(bytes) => bytes,
@@ -63,7 +69,9 @@ async fn proxy_chat_completions(
     // Check if this is a streaming request
     let is_streaming = if let Ok(body_str) = String::from_utf8(body_bytes.to_vec()) {
         if let Ok(json) = serde_json::from_str::<Value>(&body_str) {
-            json.get("stream").and_then(|v| v.as_bool()).unwrap_or(false)
+            json.get("stream")
+                .and_then(|v| v.as_bool())
+                .unwrap_or(false)
         } else {
             false
         }
@@ -72,7 +80,8 @@ async fn proxy_chat_completions(
     };
 
     // Forward the request
-    let mut req_builder = proxy_client.client
+    let mut req_builder = proxy_client
+        .client
         .post(&target_url)
         .body(body_bytes.to_vec());
 
@@ -85,8 +94,7 @@ async fn proxy_chat_completions(
 
     match req_builder.send().await {
         Ok(response) => {
-            let mut resp_builder = Response::builder()
-                .status(response.status());
+            let mut resp_builder = Response::builder().status(response.status());
 
             // Forward response headers
             for (name, value) in response.headers().iter() {
@@ -99,14 +107,12 @@ async fn proxy_chat_completions(
             if is_streaming {
                 // For streaming, we need to forward the response as-is
                 match response.bytes().await {
-                    Ok(body) => {
-                        resp_builder
-                            .header("content-type", "text/plain; charset=utf-8")
-                            .header("cache-control", "no-cache")
-                            .header("connection", "keep-alive")
-                            .body(Body::from(body))
-                            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
-                    }
+                    Ok(body) => resp_builder
+                        .header("content-type", "text/plain; charset=utf-8")
+                        .header("cache-control", "no-cache")
+                        .header("connection", "keep-alive")
+                        .body(Body::from(body))
+                        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR),
                     Err(e) => {
                         tracing::error!("Failed to read streaming response body: {}", e);
                         Err(StatusCode::INTERNAL_SERVER_ERROR)
@@ -115,11 +121,9 @@ async fn proxy_chat_completions(
             } else {
                 // For non-streaming, forward the JSON response
                 match response.bytes().await {
-                    Ok(body) => {
-                        resp_builder
-                            .body(Body::from(body))
-                            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
-                    }
+                    Ok(body) => resp_builder
+                        .body(Body::from(body))
+                        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR),
                     Err(e) => {
                         tracing::error!("Failed to read response body: {}", e);
                         Err(StatusCode::INTERNAL_SERVER_ERROR)
@@ -139,10 +143,16 @@ async fn proxy_models(
     State(proxy_client): State<ProxyClient>,
     headers: HeaderMap,
 ) -> Result<Response, StatusCode> {
-    let target_url = format!("{}/v1/models", proxy_client.config.inference_url());
-    
+    let target_url = format!(
+        "{}/v1/models",
+        proxy_client
+            .config
+            .inference_url()
+            .expect("Invalid Configuration Detected")
+    );
+
     tracing::info!("Proxying models request to: {}", target_url);
-    
+
     let mut req_builder = proxy_client.client.get(&target_url);
 
     // Forward relevant headers
@@ -154,8 +164,7 @@ async fn proxy_models(
 
     match req_builder.send().await {
         Ok(response) => {
-            let mut resp_builder = Response::builder()
-                .status(response.status());
+            let mut resp_builder = Response::builder().status(response.status());
 
             // Forward response headers
             for (name, value) in response.headers().iter() {
@@ -165,11 +174,9 @@ async fn proxy_models(
             }
 
             match response.bytes().await {
-                Ok(body) => {
-                    resp_builder
-                        .body(Body::from(body))
-                        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
-                }
+                Ok(body) => resp_builder
+                    .body(Body::from(body))
+                    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR),
                 Err(e) => {
                     tracing::error!("Failed to read models response body: {}", e);
                     Err(StatusCode::INTERNAL_SERVER_ERROR)
@@ -189,10 +196,16 @@ async fn proxy_embeddings(
     headers: HeaderMap,
     body: Body,
 ) -> Result<Response, StatusCode> {
-    let target_url = format!("{}/v1/embeddings", proxy_client.config.embeddings_url());
-    
+    let target_url = format!(
+        "{}/v1/embeddings",
+        proxy_client
+            .config
+            .embeddings_url()
+            .expect("Invalid Configuration Detected")
+    );
+
     tracing::info!("Proxying embeddings request to: {}", target_url);
-    
+
     // Extract body as bytes
     let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
         Ok(bytes) => bytes,
@@ -203,7 +216,8 @@ async fn proxy_embeddings(
     };
 
     // Forward the request
-    let mut req_builder = proxy_client.client
+    let mut req_builder = proxy_client
+        .client
         .post(&target_url)
         .body(body_bytes.to_vec());
 
@@ -216,8 +230,7 @@ async fn proxy_embeddings(
 
     match req_builder.send().await {
         Ok(response) => {
-            let mut resp_builder = Response::builder()
-                .status(response.status());
+            let mut resp_builder = Response::builder().status(response.status());
 
             // Forward response headers
             for (name, value) in response.headers().iter() {
@@ -227,11 +240,9 @@ async fn proxy_embeddings(
             }
 
             match response.bytes().await {
-                Ok(body) => {
-                    resp_builder
-                        .body(Body::from(body))
-                        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
-                }
+                Ok(body) => resp_builder
+                    .body(Body::from(body))
+                    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR),
                 Err(e) => {
                     tracing::error!("Failed to read embeddings response body: {}", e);
                     Err(StatusCode::INTERNAL_SERVER_ERROR)
@@ -250,7 +261,7 @@ fn should_forward_header(header_name: &str) -> bool {
     match header_name.to_lowercase().as_str() {
         "content-type" | "content-length" | "authorization" | "user-agent" | "accept" => true,
         "host" | "connection" | "upgrade" => false, // Don't forward connection-specific headers
-        _ => true, // Forward other headers by default
+        _ => true,                                  // Forward other headers by default
     }
 }
 
@@ -259,7 +270,7 @@ fn should_forward_response_header(header_name: &str) -> bool {
     match header_name.to_lowercase().as_str() {
         "content-type" | "content-length" | "cache-control" | "connection" => true,
         "server" | "date" => false, // Don't forward server-specific headers
-        _ => true, // Forward other headers by default
+        _ => true,                  // Forward other headers by default
     }
 }
 
@@ -290,14 +301,20 @@ mod tests {
             server_host: "127.0.0.1".to_string(),
             server_port: 8080,
             server_mode: ServerMode::HighAvailability,
-            services: Services {
-                inference_url: "http://test-inference:8080".to_string(),
-                embeddings_url: "http://test-embeddings:8080".to_string(),
-            },
+            services: Some(Services {
+                inference_url: Some("http://test-inference:8080".to_string()),
+                embeddings_url: Some("http://test-embeddings:8080".to_string()),
+            }),
         };
 
         let proxy_client = ProxyClient::new(config);
-        assert_eq!(proxy_client.config.inference_url(), "http://test-inference:8080");
-        assert_eq!(proxy_client.config.embeddings_url(), "http://test-embeddings:8080");
+        assert_eq!(
+            proxy_client.config.inference_url().unwrap().as_str(),
+            "http://test-inference:8080"
+        );
+        assert_eq!(
+            proxy_client.config.embeddings_url().unwrap().as_str(),
+            "http://test-embeddings:8080"
+        );
     }
-}
+}

crates/predict-otron-9000/src/standalone.rs

@@ -0,0 +1,19 @@
+use crate::config::ServerConfig;
+use axum::Router;
+use inference_engine::AppState;
+
+pub fn create_standalone_router(server_config: ServerConfig) -> Router {
+    // Create unified router by merging embeddings and inference routers (existing behavior)
+    let embeddings_router = embeddings_engine::create_embeddings_router();
+
+    // Create AppState with correct model configuration
+    let app_state = AppState::default();
+
+    // Get the inference router directly from the inference engine
+    let inference_router = inference_engine::create_router(app_state);
+
+    // Merge the local routers
+    Router::new()
+        .merge(embeddings_router)
+        .merge(inference_router)
+}

docs/ARCHITECTURE.md

@@ -22,7 +22,7 @@ The Predict-Otron-9000 is a comprehensive multi-service AI platform built around
 graph TB
     subgraph "Core Components"
         A[Main Server<br/>predict-otron-9000]
-        B[Inference Engine<br/>Gemma via Candle]
+        B[Inference Engine<br/>Gemma/Llama via Candle]
         C[Embeddings Engine<br/>FastEmbed]
         D[Web Frontend<br/>Leptos WASM]
     end
@@ -52,7 +52,7 @@ graph TB
 
 ## Workspace Structure
 
-The project uses a 4-crate Rust workspace with TypeScript tooling, designed for maximum flexibility in deployment configurations.
+The project uses a 7-crate Rust workspace with TypeScript tooling, designed for maximum flexibility in deployment configurations.
 
 ```mermaid
 graph TD
@@ -62,24 +62,33 @@ graph TD
         end
         
         subgraph "AI Services"
-            B[inference-engine<br/>Edition: 2021<br/>Port: 8080<br/>Candle ML]
+            B[inference-engine<br/>Edition: 2021<br/>Port: 8080<br/>Multi-model orchestrator]
+            J[gemma-runner<br/>Edition: 2021<br/>Gemma via Candle]
+            K[llama-runner<br/>Edition: 2021<br/>Llama via Candle]
             C[embeddings-engine<br/>Edition: 2024<br/>Port: 8080<br/>FastEmbed]
         end
         
         subgraph "Frontend"
             D[leptos-app<br/>Edition: 2021<br/>Port: 3000/8788<br/>WASM/SSR]
         end
+        
+        subgraph "Tooling"
+            L[helm-chart-tool<br/>Edition: 2024<br/>K8s deployment]
+        end
     end
     
     subgraph "External Tooling"
-        E[cli.ts<br/>TypeScript/Bun<br/>OpenAI SDK]
+        E[scripts/cli.ts<br/>TypeScript/Bun<br/>OpenAI SDK]
     end
     
     subgraph "Dependencies"
         A --> B
         A --> C
         A --> D
-        B -.-> F[Candle 0.9.1]
+        B --> J
+        B --> K
+        J -.-> F[Candle 0.9.1]
+        K -.-> F
         C -.-> G[FastEmbed 4.x]
         D -.-> H[Leptos 0.8.0]
         E -.-> I[OpenAI SDK 5.16+]
@@ -87,9 +96,12 @@ graph TD
     
     style A fill:#e1f5fe
     style B fill:#f3e5f5
+    style J fill:#f3e5f5
+    style K fill:#f3e5f5
     style C fill:#e8f5e8
     style D fill:#fff3e0
     style E fill:#fce4ec
+    style L fill:#fff9c4
 ```
 
 ## Deployment Configurations

scripts/build_all_platforms.sh

@@ -0,0 +1,389 @@
+#!/bin/bash
+
+# Cross-platform build script for predict-otron-9000
+# Builds all workspace crates for common platforms
+
+set -euo pipefail
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+# Configuration
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+BUILD_DIR="${PROJECT_ROOT}/build"
+TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+
+# Supported platforms
+PLATFORMS=(
+    "x86_64-unknown-linux-gnu"
+    "x86_64-pc-windows-msvc"
+    "x86_64-apple-darwin"
+    "aarch64-apple-darwin"
+    "aarch64-unknown-linux-gnu"
+)
+
+# Main binaries to build
+MAIN_BINARIES=(
+    "predict-otron-9000"
+    "embeddings-engine"
+)
+
+# Inference engine binaries (with bin feature)
+INFERENCE_BINARIES=(
+    "gemma_inference"
+    "llama_inference"
+)
+
+# Other workspace binaries
+OTHER_BINARIES=(
+    "helm-chart-tool"
+)
+
+print_header() {
+    echo -e "${BLUE}================================${NC}"
+    echo -e "${BLUE}$1${NC}"
+    echo -e "${BLUE}================================${NC}"
+}
+
+print_info() {
+    echo -e "${GREEN}[INFO]${NC} $1"
+}
+
+print_warn() {
+    echo -e "${YELLOW}[WARN]${NC} $1"
+}
+
+print_error() {
+    echo -e "${RED}[ERROR]${NC} $1"
+}
+
+check_dependencies() {
+    print_header "Checking Dependencies"
+    
+    # Check rust
+    if ! command -v cargo >/dev/null 2>&1; then
+        print_error "Rust/Cargo is not installed"
+        exit 1
+    fi
+    
+    # Check cargo-leptos for WASM frontend
+    if ! command -v cargo-leptos >/dev/null 2>&1; then
+        print_warn "cargo-leptos not found. Installing..."
+        cargo install cargo-leptos
+    fi
+    
+    print_info "All dependencies available"
+}
+
+install_targets() {
+    print_header "Installing Rust Targets"
+    
+    for platform in "${PLATFORMS[@]}"; do
+        print_info "Installing target: $platform"
+        rustup target add "$platform" || {
+            print_warn "Failed to install target $platform (may not be available on this host)"
+        }
+    done
+    
+    # Add WASM target for leptos
+    print_info "Installing wasm32-unknown-unknown target for Leptos"
+    rustup target add wasm32-unknown-unknown
+}
+
+create_build_dirs() {
+    print_header "Setting up Build Directory"
+    
+    rm -rf "$BUILD_DIR"
+    mkdir -p "$BUILD_DIR"
+    
+    for platform in "${PLATFORMS[@]}"; do
+        mkdir -p "$BUILD_DIR/$platform"
+    done
+    
+    mkdir -p "$BUILD_DIR/web"
+    print_info "Build directories created"
+}
+
+build_leptos_app() {
+    print_header "Building Leptos Web Frontend"
+    
+    cd "$PROJECT_ROOT/crates/leptos-app"
+    
+    # Build the WASM frontend
+    print_info "Building WASM frontend with cargo-leptos..."
+    cargo leptos build --release || {
+        print_error "Failed to build Leptos WASM frontend"
+        return 1
+    }
+    
+    # Copy built assets to build directory
+    if [ -d "target/site" ]; then
+        cp -r target/site/* "$BUILD_DIR/web/"
+        print_info "Leptos frontend built and copied to $BUILD_DIR/web/"
+    else
+        print_error "Leptos build output not found at target/site"
+        return 1
+    fi
+    
+    cd "$PROJECT_ROOT"
+}
+
+get_platform_features() {
+    local platform="$1"
+    local features=""
+    
+    case "$platform" in
+        *-apple-darwin)
+            # macOS uses Metal but routes to CPU for Gemma stability
+            features=""
+            ;;
+        *-unknown-linux-gnu|*-pc-windows-msvc)
+            # Linux and Windows can use CUDA if available
+            features=""
+            ;;
+        *)
+            features=""
+            ;;
+    esac
+    
+    echo "$features"
+}
+
+build_binary_for_platform() {
+    local binary_name="$1"
+    local platform="$2"
+    local package_name="$3"
+    local additional_args="$4"
+    
+    print_info "Building $binary_name for $platform"
+    
+    local features=$(get_platform_features "$platform")
+    local feature_flag=""
+    if [ -n "$features" ]; then
+        feature_flag="--features $features"
+    fi
+    
+    # Build command
+    local build_cmd="cargo build --release --target $platform --bin $binary_name"
+    
+    if [ -n "$package_name" ]; then
+        build_cmd="$build_cmd --package $package_name"
+    fi
+    
+    if [ -n "$additional_args" ]; then
+        build_cmd="$build_cmd $additional_args"
+    fi
+    
+    if [ -n "$feature_flag" ]; then
+        build_cmd="$build_cmd $feature_flag"
+    fi
+    
+    print_info "Running: $build_cmd"
+    
+    if eval "$build_cmd"; then
+        # Copy binary to build directory
+        local target_dir="target/$platform/release"
+        local binary_file="$binary_name"
+        
+        # Add .exe extension for Windows
+        if [[ "$platform" == *-pc-windows-msvc ]]; then
+            binary_file="$binary_name.exe"
+        fi
+        
+        if [ -f "$target_dir/$binary_file" ]; then
+            cp "$target_dir/$binary_file" "$BUILD_DIR/$platform/"
+            print_info "✓ $binary_name built and copied for $platform"
+        else
+            print_error "Binary not found: $target_dir/$binary_file"
+            return 1
+        fi
+    else
+        print_error "Failed to build $binary_name for $platform"
+        return 1
+    fi
+}
+
+build_for_platform() {
+    local platform="$1"
+    print_header "Building for $platform"
+    
+    local failed_builds=()
+    
+    # Build main binaries
+    for binary in "${MAIN_BINARIES[@]}"; do
+        if ! build_binary_for_platform "$binary" "$platform" "$binary" ""; then
+            failed_builds+=("$binary")
+        fi
+    done
+    
+    # Build inference engine binaries with bin feature
+    for binary in "${INFERENCE_BINARIES[@]}"; do
+        if ! build_binary_for_platform "$binary" "$platform" "inference-engine" "--features bin"; then
+            failed_builds+=("$binary")
+        fi
+    done
+    
+    # Build other workspace binaries
+    for binary in "${OTHER_BINARIES[@]}"; do
+        if ! build_binary_for_platform "$binary" "$platform" "$binary" ""; then
+            failed_builds+=("$binary")
+        fi
+    done
+    
+    if [ ${#failed_builds[@]} -eq 0 ]; then
+        print_info "✓ All binaries built successfully for $platform"
+    else
+        print_warn "Some builds failed for $platform: ${failed_builds[*]}"
+    fi
+}
+
+create_archives() {
+    print_header "Creating Release Archives"
+    
+    cd "$BUILD_DIR"
+    
+    for platform in "${PLATFORMS[@]}"; do
+        if [ -d "$platform" ] && [ -n "$(ls -A "$platform" 2>/dev/null)" ]; then
+            local archive_name="predict-otron-9000-${platform}-${TIMESTAMP}"
+            
+            print_info "Creating archive for $platform"
+            
+            # Create platform-specific directory with all files
+            mkdir -p "$archive_name"
+            cp -r "$platform"/* "$archive_name/"
+            
+            # Add web assets to each platform archive
+            if [ -d "web" ]; then
+                mkdir -p "$archive_name/web"
+                cp -r web/* "$archive_name/web/"
+            fi
+            
+            # Create README for the platform
+            cat > "$archive_name/README.txt" << EOF
+Predict-Otron-9000 - Platform: $platform
+Build Date: $(date)
+========================================
+
+Binaries included:
+$(ls -1 "$platform")
+
+Web Frontend:
+- Located in the 'web' directory
+- Serve with any static file server on port 8788 or configure your server
+
+Usage:
+1. Start the main server: ./predict-otron-9000
+2. Start embeddings service: ./embeddings-engine  
+3. Access web interface at http://localhost:8080 (served by main server)
+
+For more information, visit: https://github.com/geoffsee/predict-otron-9000
+EOF
+            
+            # Create tar.gz archive
+            tar -czf "${archive_name}.tar.gz" "$archive_name"
+            rm -rf "$archive_name"
+            
+            print_info "✓ Created ${archive_name}.tar.gz"
+        else
+            print_warn "No binaries found for $platform, skipping archive"
+        fi
+    done
+    
+    cd "$PROJECT_ROOT"
+}
+
+generate_build_report() {
+    print_header "Build Report"
+    
+    echo "Build completed at: $(date)"
+    echo "Build directory: $BUILD_DIR"
+    echo ""
+    echo "Archives created:"
+    ls -la "$BUILD_DIR"/*.tar.gz 2>/dev/null || echo "No archives created"
+    echo ""
+    echo "Platform directories:"
+    for platform in "${PLATFORMS[@]}"; do
+        if [ -d "$BUILD_DIR/$platform" ]; then
+            echo "  $platform:"
+            ls -la "$BUILD_DIR/$platform" | sed 's/^/    /'
+        fi
+    done
+    
+    if [ -d "$BUILD_DIR/web" ]; then
+        echo ""
+        echo "Web frontend assets:"
+        ls -la "$BUILD_DIR/web" | head -10 | sed 's/^/    /'
+        if [ $(ls -1 "$BUILD_DIR/web" | wc -l) -gt 10 ]; then
+            echo "    ... and $(( $(ls -1 "$BUILD_DIR/web" | wc -l) - 10 )) more files"
+        fi
+    fi
+}
+
+main() {
+    print_header "Predict-Otron-9000 Cross-Platform Build Script"
+    
+    cd "$PROJECT_ROOT"
+    
+    check_dependencies
+    install_targets
+    create_build_dirs
+    
+    # Build Leptos web frontend first
+    build_leptos_app
+    
+    # Build for each platform
+    for platform in "${PLATFORMS[@]}"; do
+        build_for_platform "$platform"
+    done
+    
+    create_archives
+    generate_build_report
+    
+    print_header "Build Complete!"
+    print_info "All artifacts are available in: $BUILD_DIR"
+}
+
+# Handle command line arguments
+case "${1:-}" in
+    --help|-h)
+        echo "Usage: $0 [options]"
+        echo ""
+        echo "Cross-platform build script for predict-otron-9000"
+        echo ""
+        echo "Options:"
+        echo "  --help, -h          Show this help message"
+        echo "  --platforms         Show supported platforms"
+        echo "  --clean             Clean build directory before building"
+        echo ""
+        echo "Supported platforms:"
+        for platform in "${PLATFORMS[@]}"; do
+            echo "  - $platform"
+        done
+        echo ""
+        echo "Prerequisites:"
+        echo "  - Rust toolchain with rustup"
+        echo "  - cargo-leptos (will be installed if missing)"
+        echo "  - Platform-specific toolchains for cross-compilation"
+        echo ""
+        exit 0
+        ;;
+    --platforms)
+        echo "Supported platforms:"
+        for platform in "${PLATFORMS[@]}"; do
+            echo "  - $platform"
+        done
+        exit 0
+        ;;
+    --clean)
+        print_info "Cleaning build directory..."
+        rm -rf "$BUILD_DIR"
+        print_info "Build directory cleaned"
+        ;;
+esac
+
+main "$@"

scripts/build_cli.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env sh
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+
+TEMP_DIR="$SCRIPT_DIR/temp"
+
+mkdir -p "$TEMP_DIR"
+
+cp "$SCRIPT_DIR/cli.ts" "$TEMP_DIR/cli.ts"
+cp "$SCRIPT_DIR/../package.json" "$TEMP_DIR/package.json"
+
+(
+cd "$TEMP_DIR"
+bun i
+bun build ./cli.ts --compile --outfile "$SCRIPT_DIR/cli"
+)
+
+rm -rf "$TEMP_DIR"