File size: 7,785 Bytes
5973eb4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | //! Minimal dependency-free HTTP/1.1 API for token-ID-based Falcon-H1 inference.
//! For localhost/Colab tunnel demo only; deploy behind a reverse proxy for production.
use falcon_h1_rust::{load_pack,synthetic_weights,Config,Model};
use std::io::{Read,Write};
use std::net::{TcpListener,TcpStream};
use std::sync::Arc;
use std::time::Duration;
fn extract_array(body:&str,key:&str)->Result<Vec<usize>,String>{
let needle=format!("\"{key}\"");
let idx=body.find(&needle).ok_or_else(||format!("missing {key}"))?;
let tail=&body[idx+needle.len()..];
let colon=tail.find(':').ok_or("missing :")?;
let tail=&tail[colon+1..];
let start=tail.find('[').ok_or("missing [")?;
let tail=&tail[start+1..];
let end=tail.find(']').ok_or("missing ]")?;
let data=tail[..end].trim();
if data.is_empty(){return Err("empty input_ids".into())}
let ids=data.split(',').map(|x|x.trim().parse::<usize>().map_err(|_|"invalid token ID".to_string())).collect::<Result<Vec<_>,_>>()?;
if ids.len()>2048{return Err("sequence exceeds API limit (2048 tokens)".into())}
Ok(ids)
}
fn extract_number(body:&str,key:&str,default:usize)->Result<usize,String>{
let needle=format!("\"{key}\"");
if let Some(idx)=body.find(&needle){
let tail=&body[idx+needle.len()..];
let colon=tail.find(':').ok_or("missing :")?;
let num=tail[colon+1..].trim_start().chars().take_while(|c|c.is_ascii_digit()).collect::<String>();
if num.is_empty(){return Err(format!("invalid {key}"))}
num.parse::<usize>().map_err(|_|format!("invalid {key}"))
} else {Ok(default)}
}
fn response(mut stream:TcpStream,code:&str,payload:&str){
let msg=format!("HTTP/1.1 {code}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\nAccess-Control-Allow-Origin: *\r\nAccess-Control-Allow-Methods: GET, POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type\r\n\r\n{}",payload.len(),payload);
let _=stream.write_all(msg.as_bytes());
}
fn json_err(e:&str)->String{format!("{{\"error\":\"{}\"}}",e.replace('\\',"\\\\").replace('"',"\\\"").replace('\n'," "))}
enum Backend {
Cpu(Model),
#[cfg(feature="candle")]
Candle(falcon_h1_rust::candle_backend::CandleModel),
}
impl Backend {
fn config(&self)->&Config { match self {
Self::Cpu(m)=>&m.config,
#[cfg(feature="candle")]
Self::Candle(m)=>&m.config,
}}
fn label(&self)->&'static str {match self {
Self::Cpu(_)=>"cpu-reference",
#[cfg(feature="candle")]
Self::Candle(m)=>m.device_name(),
}}
fn logits(&self,ids:&[usize])->Result<Vec<f32>,String>{match self {
Self::Cpu(m)=>m.logits(ids),
#[cfg(feature="candle")]
Self::Candle(m)=>m.logits(ids),
}}
fn generate(&self,ids:&[usize],n:usize)->Result<Vec<usize>,String>{match self {
Self::Cpu(m)=>m.generate(ids,n),
#[cfg(feature="candle")]
Self::Candle(m)=>m.generate(ids,n),
}}
}
fn handle(mut stream:TcpStream,model:Arc<Backend>){
let _=stream.set_read_timeout(Some(Duration::from_secs(15)));
let _=stream.set_write_timeout(Some(Duration::from_secs(15)));
let mut buf=vec![0u8;16384];let mut len=0usize;let mut expected=None;
loop {
if len==buf.len(){response(stream,"413 Payload Too Large",&json_err("request too large"));return}
match stream.read(&mut buf[len..]) {
Ok(0)=>break,
Ok(n)=>len+=n,
Err(_)=>return,
}
if let Some(hend)=buf[..len].windows(4).position(|w|w==b"\r\n\r\n") {
if expected.is_none(){
let headers=String::from_utf8_lossy(&buf[..hend]);
let size=headers.lines().filter_map(|l|l.split_once(':')).find(|(k,_)|k.eq_ignore_ascii_case("content-length"))
.and_then(|(_,v)|v.trim().parse::<usize>().ok()).unwrap_or(0);
expected=Some(hend+4+size);
}
if len>=expected.unwrap(){break}
}
}
let raw=String::from_utf8_lossy(&buf[..len]);
let first=raw.lines().next().unwrap_or("");
let (method,path) = {let mut p=first.split_whitespace();(p.next().unwrap_or(""),p.next().unwrap_or(""))};
if method=="OPTIONS" {response(stream,"204 No Content","");return}
if method=="GET" && path=="/health" {response(stream,"200 OK",&format!("{{\"status\":\"ok\",\"backend\":\"{}\"}}",model.label()));return}
if method=="GET" && path=="/v1/model" {
let c=model.config();
let msg=format!("{{\"architecture\":\"falcon_h1\",\"backend\":\"{}\",\"layers\":{},\"hidden\":{},\"attention_heads\":{},\"kv_heads\":{},\"mamba_heads\":{},\"vocab_size\":{}}}",model.label(),c.layers,c.hidden,c.heads,c.kv_heads,c.mamba_heads,c.vocab);
response(stream,"200 OK",&msg);return
}
if method!="POST" || (path!="/v1/logits" && path!="/v1/generate") {
response(stream,"404 Not Found",&json_err("unknown route"));return
}
let body=match raw.split_once("\r\n\r\n") {Some((_,v))=>v,None=>{response(stream,"400 Bad Request",&json_err("no body"));return}};
let result = (||->Result<String,String>{
let ids=extract_array(body,"input_ids")?;
if path=="/v1/logits" {
let y=model.logits(&ids)?;
let text=y.iter().map(|x|format!("{x:.7}")).collect::<Vec<_>>().join(",");
Ok(format!("{{\"logits\":[{text}],\"vocab_size\":{}}}",y.len()))
} else {
let count=extract_number(body,"max_new_tokens",1)?;
let seq=model.generate(&ids,count)?;
Ok(format!("{{\"output_ids\":[{}],\"generated_tokens\":{count}}}",seq.iter().map(|x|x.to_string()).collect::<Vec<_>>().join(",")))
}
})();
match result{Ok(msg)=>response(stream,"200 OK",&msg),Err(e)=>response(stream,"400 Bad Request",&json_err(&e))}
}
fn main()->Result<(),Box<dyn std::error::Error>>{
let args:Vec<String>=std::env::args().collect();
let demo=args.iter().any(|x|x=="--demo");
let bind=args.iter().position(|x|x=="--bind").and_then(|i|args.get(i+1)).map(|s|s.as_str()).unwrap_or("127.0.0.1:8080");
let c=if demo{Config::demo()}else{Config::falcon_h1_05b()};
let weights=if demo{synthetic_weights(&c)}else{
let path=args.iter().position(|x|x=="--checkpoint").and_then(|i|args.get(i+1)).ok_or("use --checkpoint path/to/fh1.f32pack or --demo")?;
load_pack(path)?
};
let device=args.iter().position(|x|x=="--device").and_then(|i|args.get(i+1)).map(|s|s.as_str()).unwrap_or("cpu");
let model=match device {
"cpu"=>Backend::Cpu(Model::new(c,weights)?),
#[cfg(feature="candle")]
"cpu-candle"=>Backend::Candle(falcon_h1_rust::candle_backend::CandleModel::new(c,weights,candle_core::Device::Cpu)?),
#[cfg(feature="cuda")]
"cuda"=>Backend::Candle(falcon_h1_rust::candle_backend::CandleModel::new(c,weights,candle_core::Device::new_cuda(0)?)?),
_=>return Err(format!("unsupported --device {device}; build with --features cuda for CUDA").into()),
};
let model=Arc::new(model);
let listener=TcpListener::bind(bind)?;
eprintln!("Falcon-H1 Rust API listening on http://{bind} (demo={demo}, backend={})",model.label());
for request in listener.incoming(){
if let Ok(stream)=request {let shared=Arc::clone(&model);
std::thread::spawn(move||handle(stream,shared));}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test] fn parse_ids(){ assert_eq!(extract_array(r#"{"input_ids":[1,2,3]}"#,"input_ids").unwrap(),vec![1,2,3]); }
#[test] fn reject_ids(){assert!(extract_array(r#"{"input_ids":[1,"x"]}"#,"input_ids").is_err());}
#[test] fn parse_count(){ assert_eq!(extract_number(r#"{"max_new_tokens":3}"#,"max_new_tokens",1).unwrap(),3); }
}
|