| |
| |
| 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); } |
| } |
|
|
