Rust AI 与大模型集成：Candle、llm.rs 与 Python FFI

use candle_core::{Device, Tensor, DType};

let device = Device::new_cuda(0)?; // 或 Device::Cpu, Device::Metal
let a = Tensor::randn(0f32, 1., (2, 3), &device)?;
let b = Tensor::randn(0f32, 1., (3, 4), &device)?;
let c = a.matmul(&b)?;
println!("{}", c);

use candle_transformers::models::bert::{BertModel, Config};
use tokenizers::Tokenizer;

let device = Device::Cpu;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_path], DType::F32, &device)? };
let config = Config::v-base();
let model = BertModel::load(vb, &config)?;

let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(E::msg)?;
let tokens = tokenizer.encode("Hello, world!", true).map_err(E::msg)?;
let token_ids = Tensor::new(tokens.get_ids(), &device)?.unsqueeze(0)?;

let embeddings = model.forward(&token_ids)?;
println!("embeddings shape: {:?}", embeddings.shape());

use candle_transformers::models::llama::{Llama, Config};

let mut model = Llama::load(vb, &config)?;
let mut logits_processor = LogitsProcessor::new(299792458, None, None);

let mut index_pos = 0;
let mut tokens = prompt_tokens;

for index in 0..sample_len {
    let context_size = if index > 0 { 1 } else { tokens.len() };
    let ctxt = &tokens[tokens.len().saturating_sub(context_size)..];
    let input = Tensor::new(ctxt, &device)?.unsqueeze(0)?;
    let logits = model.forward(&input, index_pos)?;
    let logits = logits.squeeze(0)?.to_dtype(DType::F32)?;
    let next_token = logits_processor.sample(&logits)?;
    tokens.push(next_token);
    index_pos += ctxt.len();
}

use candle_core::quantized::gguf_file;

let mut file = std::fs::File::open("llama-7b-q4_0.gguf")?;
let content = gguf_file::Content::read(&mut file)?;
let model = quantized_llama::Model::from_gguf(content, &mut file, &device)?;
// Q4_0 量化将模型从 13GB 压缩至 ~4GB，推理速度提升 2-3x

use llm::models::Llama;
use llm::Model;

let model = llm::load::<Llama>(
    &path,
    llm::TokenizerSource::Embedded,
    Default::default(),
    llm::load_progress_callback_stdout,
)?;

let mut session = model.start_session(Default::default());

session.infer(
    &model,
    &mut rand::thread_rng(),
    &llm::InferenceRequest {
        prompt: "Rust is a programming language".into(),
        ..Default::default()
    },
    &mut Default::default(),
    |token| {
        print!("{}", token);
        Ok(llm::InferenceFeedback::Continue)
    },
)?;

fn chat_prompt(system: &str, history: &[(String, String)], user: &str) -> String {
    let mut prompt = format!("<|system|>\n{}\n", system);
    for (u, a) in history {
        prompt.push_str(&format!("<|user|>\n{}\n<|assistant|>\n{}\n", u, a));
    }
    prompt.push_str(&format!("<|user|>\n{}\n<|assistant|>\n", user));
    prompt
}

use ort::{Environment, SessionBuilder, Value};

let environment = Environment::builder()
    .with_name("my_model")
    .build()?
    .into_arc();

let session = SessionBuilder::new(&environment)?
    .with_model_from_file("model.onnx")?;

// 准备输入
let input = vec![1.0f32; 3 * 224 * 224];
let input_tensor = Value::from_array(([1, 3, 224, 224], input.into_boxed_slice()))?;

// 推理
let outputs = session.run(vec![input_tensor])?;
let predictions: ndarray::ArrayViewD<f32> = outputs[0].try_extract()?;

use pyo3::prelude::*;

/// 高性能 Rust 函数暴露给 Python
#[pyfunction]
fn fast_tokenize(text: &str) -> Vec<String> {
    text.split_whitespace()
        .map(|s| s.to_lowercase())
        .collect()
}

#[pymodule]
fn my_rust_tokenizer(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(fast_tokenize, m)?)?;
    Ok(())
}

# Python 调用
import my_rust_tokenizer

tokens = my_rust_tokenizer.fast_tokenize("Hello World FROM Rust")
# ['hello', 'world', 'from', 'rust']

use pyo3::prelude::*;
use pyo3::types::PyDict;

Python::with_gil(|py| {
    let transformers = py.import("transformers")?;
    let tokenizer = transformers
        .getattr("AutoTokenizer")?
        .call_method("from_pretrained", ("bert-base-chinese",), None)?;
    
    let result = tokenizer.call_method("encode", ("你好世界",), None)?;
    let ids: Vec<u32> = result.extract()?;
    println!("{:?}", ids);
    Ok(())
})

# 安装 maturin
pip install maturin

# 开发模式（快速迭代）
maturin develop

# 构建 wheel 分发
maturin build --release

# 发布到 PyPI
maturin publish

use qdrant_client::prelude::*;
use qdrant_client::qdrant::{CreateCollection, VectorParams, Distance};

let client = QdrantClient::from_url("http://localhost:6334").build()?;

// 创建集合
client.create_collection(&CreateCollection {
    collection_name: "documents".to_string(),
    vectors_config: Some(VectorParams {
        size: 768,
        distance: Distance::Cosine.into(),
        ..Default::default()
    }.into()),
    ..Default::default()
}).await?;

use qdrant_client::qdrant::{PointStruct, UpsertPointsBuilder};

// 插入向量
let points = vec![PointStruct::new(
    1,                              // point id
    vec![0.05, 0.61, 0.76, /* ... 768 dims */],
    json!({"text": "Rust programming guide"}),
)];

client.upsert_points(UpsertPointsBuilder::new("documents", points)).await?;

// 相似度搜索
let search_result = client.search_points(&SearchPoints {
    collection_name: "documents".to_string(),
    vector: query_embedding,
    limit: 5,
    with_payload: Some(true.into()),
    ..Default::default()
}).await?;

use candle_transformers::models::bert::BertModel;

struct RAGSystem {
    embedder: BertModel,
    qdrant: QdrantClient,
    llm: Llama,
}

impl RAGSystem {
    async fn query(&self, question: &str) -> Result<String> {
        // 1. 将问题转为向量
        let query_embedding = self.embed(question).await?;
        
        // 2. 检索相关文档
        let docs = self.qdrant.search_points(&SearchPoints {
            collection_name: "docs".to_string(),
            vector: query_embedding,
            limit: 3,
            ..Default::default()
        }).await?;
        
        // 3. 构建上下文提示
        let context = docs.result.iter()
            .map(|r| r.payload["text"].as_str().unwrap())
            .collect::<Vec<_>>()
            .join("\n");
        
        let prompt = format!(
            "基于以下上下文回答问题：\n{}\n\n问题：{}\n回答：",
            context, question
        );
        
        // 4. 生成回答
        self.generate(&prompt).await
    }
}

fn chunk_text(text: &str, chunk_size: usize, overlap: usize) -> Vec<String> {
    text.chars()
        .collect::<Vec<_>>()
        .windows(chunk_size)
        .step_by(chunk_size - overlap)
        .map(|w| w.iter().collect())
        .collect()
}

// 批量编码提升 GPU 利用率
let batch_size = 32;
for chunk in inputs.chunks(batch_size) {
    let batch_tensor = Tensor::stack(
        &chunk.iter().map(|t| tokenize(t)).collect::<Result<Vec<_>>>()?,
        0
    )?;
    let outputs = model.forward(&batch_tensor)?;
    // 处理 outputs...
}

use tracing::{info, instrument};

#[instrument]
async fn inference_handler(request: InferenceRequest) -> Result<InferenceResponse> {
    let start = Instant::now();
    let result = model.generate(&request.prompt).await?;
    info!(
        latency_ms = start.elapsed().as_millis(),
        tokens_generated = result.tokens.len(),
        "inference completed"
    );
    Ok(result)
}

// Cargo.toml
// [dependencies]
// candle-core = { version = "0.3", features = ["cuda"] }
// candle-nn = "0.3"
// wasm-bindgen = "0.2"
// wasm-bindgen-futures = "0.4"
// js-sys = "0.3"
// web-sys = { version = "0.3", features = ["console"] }

#![cfg(target_arch = "wasm32")]

use candle_core::{Device, Tensor};
use candle_nn::VarBuilder;
use wasm_bindgen::prelude::*;

#[wasm_bindgen]
pub struct TextGenerator {
    model: candle_transformers::models::llama::Llama,
    tokenizer: tokenizers::Tokenizer,
}

#[wasm_bindgen]
impl TextGenerator {
    pub async fn from_weights(weights_url: &str) -> Result<TextGenerator, JsValue> {
        // 从 URL 加载量化模型权重
        let weights = fetch_weights(weights_url).await?;
        let device = Device::Cpu;
        let vb = VarBuilder::from_safetensors(vec![weights], candle_core::DType::F32, &device)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;

        let config = candle_transformers::models::llama::Config::config_7b();
        let model = candle_transformers::models::llama::Llama::load(vb, &config)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;

        let tokenizer = tokenizers::Tokenizer::from_bytes(
            include_bytes!("../tokenizer.json")
        ).map_err(|e| JsValue::from_str(&e.to_string()))?;

        Ok(TextGenerator { model, tokenizer })
    }

    pub fn generate(&self, prompt: &str, max_tokens: usize) -> Result<String, JsValue> {
        // 推理逻辑（同步，在主线程外使用 Web Worker）
        let tokens = self.tokenizer.encode(prompt, true)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        // ... 自回归生成
        Ok("generated text".to_string())
    }
}

# .github/workflows/mlops-rust.yml
name: MLOps Rust Inference Pipeline

on:
  push:
    paths:
      - 'models/**'
      - 'src/**'
      - '.github/workflows/mlops-rust.yml'

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Install Rust
        uses: dtolnay/rust@stable

      - name: Cache dependencies
        uses: Swatinem/rust-cache@v2

      - name: Run unit tests
        run: cargo test --release

      - name: Run model integration tests
        run: cargo test --release --test model_integration

  benchmark:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: dtolnay/rust@stable
      - name: Benchmark inference latency
        run: |
          cargo bench --bench inference > benchmark.txt
          cat benchmark.txt
      - name: Upload benchmark results
        uses: actions/upload-artifact@v4
        with:
          name: benchmark-results
          path: benchmark.txt

  build:
    needs: [test, benchmark]
    strategy:
      matrix:
        target:
          - x86_64-unknown-linux-gnu
          - x86_64-unknown-linux-musl
          - aarch64-unknown-linux-gnu
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: dtolnay/rust@stable
      - name: Install cross-compilation tools
        run: cargo install cross
      - name: Build release binary
        run: cross build --release --target ${{ matrix.target }}
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
          name: inference-server-${{ matrix.target }}
          path: target/${{ matrix.target }}/release/inference-server

  docker:
    needs: build
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - name: Build and push Docker image
        run: |
          docker build -t ghcr.io/myorg/inference-server:${{ github.sha }} .
          docker push ghcr.io/myorg/inference-server:${{ github.sha }}
          docker tag ghcr.io/myorg/inference-server:${{ github.sha }} \
                     ghcr.io/myorg/inference-server:latest
          docker push ghcr.io/myorg/inference-server:latest

  deploy:
    needs: docker
    if: github.ref == 'refs/heads/main'
    runs-on: ubuntu-latest
    steps:
      - name: Deploy to K8s
        run: |
          kubectl set image deployment/inference-server \
            inference-server=ghcr.io/myorg/inference-server:${{ github.sha }}
          kubectl rollout status deployment/inference-server

use candle_core::safetensors::load;

// ✅ Safe: safetensors 格式，无代码执行风险
let tensors = load("model.safetensors", &device)?;

// ❌ 危险: PyTorch pickle 可能执行任意代码
// let tensors = candle_core::pickle::load("model.pt", &device)?;

const FORBIDDEN_PATTERNS: &[&str] = &[
    "ignore previous instructions",
    "system prompt",
    "DAN mode",
    "jailbreak",
];

fn sanitize_input(input: &str) -> Result<String, anyhow::Error> {
    let lower = input.to_lowercase();
    for pattern in FORBIDDEN_PATTERNS {
        if lower.contains(pattern) {
            bail!("Input contains potentially malicious content")
        }
    }

    // 限制输入长度
    if input.chars().count() > 4096 {
        bail!("Input exceeds maximum length of 4096 characters")
    }

    Ok(input.to_string())
}

use std::time::{Duration, Instant};
use tokio::time::timeout;

async fn safe_inference(
    model: &mut Llama,
    prompt: &str,
    max_tokens: usize,
    max_duration: Duration,
) -> Result<String, InferenceError> {
    let start = Instant::now();

    let result = timeout(max_duration, async {
        model.generate(prompt, max_tokens).await
    }).await;

    match result {
        Ok(Ok(text)) => {
            tracing::info!(
                latency_ms = start.elapsed().as_millis(),
                tokens_generated = text.split_whitespace().count(),
                "inference completed"
            );
            Ok(text)
        }
        Ok(Err(e)) => Err(InferenceError::ModelError(e)),
        Err(_) => Err(InferenceError::Timeout),
    }
}