File size: 890 Bytes
9db8ced
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import type { TextGenerationStreamOutput } from "@huggingface/inference";
import type OpenAI from "openai";
import type { Stream } from "openai/streaming";

/**
 * Transform a stream of OpenAI.Completions.Completion into a stream of TextGenerationStreamOutput
 */
export async function* openAICompletionToTextGenerationStream(
	completionStream: Stream<OpenAI.Completions.Completion>
) {
	let generatedText = "";
	let tokenId = 0;
	for await (const completion of completionStream) {
		const { choices } = completion;
		const text = choices[0]?.text ?? "";
		const last = choices[0]?.finish_reason === "stop";
		if (text) {
			generatedText = generatedText + text;
		}
		const output: TextGenerationStreamOutput = {
			token: {
				id: tokenId++,
				text,
				logprob: 0,
				special: false,
			},
			generated_text: last ? generatedText : null,
			details: null,
		};
		yield output;
	}
}