3
3
import org .jetbrains .annotations .NotNull ;
4
4
5
5
import java .nio .charset .StandardCharsets ;
6
+ import java .util .Collections ;
6
7
import java .util .Map ;
7
8
8
9
/**
@@ -24,16 +25,72 @@ public final class LlamaOutput {
24
25
@ NotNull
25
26
public final Map <String , Float > probabilities ;
26
27
28
+ /**
29
+ * Indicates whether this is the final output in the sequence
30
+ */
27
31
final boolean stop ;
28
32
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
+ */
29
40
LlamaOutput (byte [] generated , @ NotNull Map <String , Float > probabilities , boolean stop ) {
30
41
this .text = new String (generated , StandardCharsets .UTF_8 );
31
42
this .probabilities = probabilities ;
32
43
this .stop = stop ;
33
44
}
34
45
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
+
35
92
@ Override
36
93
public String toString () {
37
94
return text ;
38
95
}
39
- }
96
+ }
0 commit comments