47 lines
1.4 KiB
TypeScript
47 lines
1.4 KiB
TypeScript
import {
|
|
type AIStreamCallbacksAndOptions,
|
|
createCallbacksTransformer,
|
|
trimStartOfStreamHelper,
|
|
} from './ai-stream';
|
|
import { createStreamDataTransformer } from './stream-data';
|
|
|
|
function createParser(res: AsyncGenerator<any>) {
|
|
const trimStartOfStream = trimStartOfStreamHelper();
|
|
return new ReadableStream<string>({
|
|
async pull(controller): Promise<void> {
|
|
const { value, done } = await res.next();
|
|
|
|
if (done) {
|
|
controller.close();
|
|
return;
|
|
}
|
|
|
|
const text = trimStartOfStream(value.token?.text ?? '');
|
|
if (!text) return;
|
|
|
|
// some HF models return generated_text instead of a real ending token
|
|
if (value.generated_text != null && value.generated_text.length > 0) {
|
|
return;
|
|
}
|
|
|
|
// <|endoftext|> is for https://huggingface.co/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5
|
|
// <|end|> is for https://huggingface.co/HuggingFaceH4/starchat-beta
|
|
// </s> is also often last token in the stream depending on the model
|
|
if (text === '</s>' || text === '<|endoftext|>' || text === '<|end|>') {
|
|
return;
|
|
}
|
|
|
|
controller.enqueue(text);
|
|
},
|
|
});
|
|
}
|
|
|
|
export function HuggingFaceStream(
|
|
res: AsyncGenerator<any>,
|
|
callbacks?: AIStreamCallbacksAndOptions,
|
|
): ReadableStream {
|
|
return createParser(res)
|
|
.pipeThrough(createCallbacksTransformer(callbacks))
|
|
.pipeThrough(createStreamDataTransformer());
|
|
}
|