Skip to content

Commit d09c3fc

Browse files
author
Vaijanath Rao
committed
support for stream and regular
1 parent 64f80c1 commit d09c3fc

File tree

1 file changed

+58
-1
lines changed

1 file changed

+58
-1
lines changed

src/main/java/de/kherud/llama/LlamaOutput.java

+58-1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import org.jetbrains.annotations.NotNull;
44

55
import java.nio.charset.StandardCharsets;
6+
import java.util.Collections;
67
import java.util.Map;
78

89
/**
@@ -24,16 +25,72 @@ public final class LlamaOutput {
2425
@NotNull
2526
public final Map<String, Float> probabilities;
2627

28+
/**
29+
* Indicates whether this is the final output in the sequence
30+
*/
2731
final boolean stop;
2832

33+
/**
34+
* Create a new LlamaOutput from raw bytes with probability information.
35+
*
36+
* @param generated The raw bytes of generated text
37+
* @param probabilities Mapping of tokens to their probabilities
38+
* @param stop Whether this is the final output
39+
*/
2940
LlamaOutput(byte[] generated, @NotNull Map<String, Float> probabilities, boolean stop) {
3041
this.text = new String(generated, StandardCharsets.UTF_8);
3142
this.probabilities = probabilities;
3243
this.stop = stop;
3344
}
3445

46+
/**
47+
* Create a new LlamaOutput from a string with probability information.
48+
*
49+
* @param text The generated text as a string
50+
* @param probabilities Mapping of tokens to their probabilities
51+
* @param stop Whether this is the final output
52+
*/
53+
public LlamaOutput(@NotNull String text, @NotNull Map<String, Float> probabilities, boolean stop) {
54+
this.text = text;
55+
this.probabilities = probabilities;
56+
this.stop = stop;
57+
}
58+
59+
/**
60+
* Create a new LlamaOutput from a string with no probability information.
61+
* This is useful for chat completions where token probabilities might not be available.
62+
*
63+
* @param text The generated text as a string
64+
* @param stop Whether this is the final output
65+
*/
66+
public LlamaOutput(@NotNull String text, boolean stop) {
67+
this.text = text;
68+
this.probabilities = Collections.emptyMap();
69+
this.stop = stop;
70+
}
71+
72+
/**
73+
* Create a new LlamaOutput by wrapping a StreamingOutput object.
74+
*
75+
* @param output The streaming output to wrap
76+
*/
77+
public LlamaOutput(@NotNull StreamingOutput output) {
78+
this.text = output.text;
79+
this.probabilities = output.probabilities != null ? output.probabilities : Collections.emptyMap();
80+
this.stop = output.isFinal;
81+
}
82+
83+
/**
84+
* Returns whether this is the final output in a sequence
85+
*
86+
* @return true if this is the final output, false otherwise
87+
*/
88+
public boolean isStop() {
89+
return stop;
90+
}
91+
3592
@Override
3693
public String toString() {
3794
return text;
3895
}
39-
}
96+
}

0 commit comments

Comments
 (0)