Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitca7e4ad

Browse files
committed
make bindings vllm::outputs
1 parent746953e commitca7e4ad

File tree

4 files changed

+67
-15
lines changed

4 files changed

+67
-15
lines changed

‎pgml-extension/src/api.rs‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -653,7 +653,19 @@ fn transform(mut task: Value, args: Value, inputs: Vec<&str>) -> anyhow::Result<
653653
builder.build().unwrap()
654654
});
655655

656-
Ok(json!(llm.generate(&inputs,None)?))
656+
let outputs = llm
657+
.generate(&inputs,None)?
658+
.iter()
659+
.map(|o|{
660+
o.outputs()
661+
.unwrap()
662+
.iter()
663+
.map(|o| o.text().unwrap())
664+
.collect::<Vec<_>>()
665+
})
666+
.collect::<Vec<Vec<_>>>();
667+
668+
Ok(json!(outputs))
657669
}else{
658670
ifletSome(map) = task.as_object_mut(){
659671
// pop backend keyword, if present

‎pgml-extension/src/bindings/vllm/llm.rs‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use pyo3::{prelude::*, types::PyDict};
22
use serde_json::Value;
33

4-
usesuper::SamplingParams;
4+
usesuper::{RequestOutput,SamplingParams};
55

66
pubstructLLMBuilder{
77
model:String,
@@ -170,27 +170,18 @@ impl LLM {
170170
&self,
171171
prompts:&[&str],
172172
params:Option<&SamplingParams>,
173-
) ->PyResult<Vec<String>>{
173+
) ->PyResult<Vec<RequestOutput>>{
174174
let prompts:Vec<_> = prompts.iter().map(|s| s.to_string()).collect();
175175

176176
Python::with_gil(|py|{
177177
let kwargs =PyDict::new(py);
178178
kwargs.set_item("prompts", prompts)?;
179179
kwargs.set_item("sampling_params", params)?;
180180

181-
let outputs:Vec<PyObject> =self
182-
.inner
181+
self.inner
183182
.getattr(py,"generate")?
184183
.call(py,(),Some(kwargs))?
185-
.extract(py)?;
186-
187-
outputs
188-
.iter()
189-
.map(|output| ->PyResult<String>{
190-
let outputs:Vec<PyObject> = output.getattr(py,"outputs")?.extract(py)?;
191-
outputs.first().unwrap().getattr(py,"text")?.extract(py)
192-
})
193-
.collect::<PyResult<Vec<_>>>()
184+
.extract(py)
194185
})
195186
}
196187
}
@@ -238,7 +229,7 @@ impl TryFrom<Value> for LLMBuilder {
238229

239230
#[cfg(test)]
240231
mod tests{
241-
usecrate::SamplingParamsBuilder;
232+
usecrate::bindings::vllm::SamplingParamsBuilder;
242233

243234
usesuper::*;
244235

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
//! Rust bindings to the Python package `vllm`.
22
33
mod llm;
4+
mod outputs;
45
mod params;
56

67
pubuse llm::*;
8+
pubuse outputs::*;
79
pubuse params::*;
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
use pyo3::prelude::*;
2+
3+
#[derive(Clone)]
4+
pubstructRequestOutput{
5+
inner:PyObject,
6+
}
7+
8+
#[derive(Clone)]
9+
pubstructCompletionOutput{
10+
inner:PyObject,
11+
}
12+
13+
implRequestOutput{
14+
pubfnprompt(&self) ->PyResult<String>{
15+
Python::with_gil(|py|self.inner.getattr(py,"prompt")?.extract(py))
16+
}
17+
18+
pubfnoutputs(&self) ->PyResult<Vec<CompletionOutput>>{
19+
Python::with_gil(|py|self.inner.getattr(py,"outputs")?.extract(py))
20+
}
21+
}
22+
23+
implCompletionOutput{
24+
pubfnfinished(&self) ->PyResult<bool>{
25+
Python::with_gil(|py|self.inner.getattr(py,"finished")?.call0(py)?.extract(py))
26+
}
27+
28+
pubfntext(&self) ->PyResult<String>{
29+
Python::with_gil(|py|self.inner.getattr(py,"text")?.extract(py))
30+
}
31+
}
32+
33+
impl<'source>FromPyObject<'source>forRequestOutput{
34+
fnextract(ob:&'sourcePyAny) ->PyResult<Self>{
35+
Ok(Self{
36+
inner: ob.extract()?,
37+
})
38+
}
39+
}
40+
41+
impl<'source>FromPyObject<'source>forCompletionOutput{
42+
fnextract(ob:&'sourcePyAny) ->PyResult<Self>{
43+
Ok(Self{
44+
inner: ob.extract()?,
45+
})
46+
}
47+
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp