Request for Chinese Language Support in Whisper Model about djl HOT 13 CLOSED

        Criteria<Audio, String> criteria =
                Criteria.builder()
                        .setTypes(Audio.class, String.class)
                        .optModelUrls("file:build/example/whisper-large-v2.pt")
                        .optEngine("PyTorch")
                        .optDevice(Device.cpu())
                        .optTranslator(new CustWhisperTranslator())
                        .build();
        whisperModel = criteria.loadModel();
        predictor = whisperModel.newPredictor();

public class CustWhisperTranslator implements NoBatchifyTranslator<Audio, String> {

    private static final Map<Character, Byte> BYTES_DECODER = bpeDecoder();

    private List<AudioProcessor> processors;
    private Vocabulary vocabulary;

    public CustWhisperTranslator() {
        processors = new ArrayList<>();
    }

    /** {@inheritDoc} */
    @Override
    public void prepare(TranslatorContext ctx) throws IOException {
        Path path = ctx.getModel().getModelPath();
        Path melFile = path.resolve("mel_80_filters.npz");

        processors.add(new PadOrTrim(480000));
        // Use model's NDManager
        NDManager modelManager = ctx.getModel().getNDManager();
        processors.add(LogMelSpectrogram.newInstance(melFile, 80, modelManager));

        Map<String, Integer> vocab;
        Map<String, Integer> added;
        Type type = new TypeToken<Map<String, Integer>>() {}.getType();
        try (Reader reader = Files.newBufferedReader(path.resolve("vocab.json"))) {
            vocab = JsonUtils.GSON.fromJson(reader, type);
        }
        try (Reader reader = Files.newBufferedReader(path.resolve("added_tokens.json"))) {
            added = JsonUtils.GSON.fromJson(reader, type);
        }
        String[] result = new String[vocab.size() + added.size()];
        vocab.forEach((key, value) -> result[value] = key);
        added.forEach((key, value) -> result[value] = key);
        vocabulary = new DefaultVocabulary(Arrays.asList(result));
    }

    /** {@inheritDoc} */
    @Override
    public NDList processInput(TranslatorContext ctx, Audio input) throws Exception {
        NDArray samples = ctx.getNDManager().create(input.getData());
        for (AudioProcessor processor : processors) {
            samples = processor.extractFeatures(samples.getManager(), samples);
        }
        samples = samples.expandDims(0);
        NDArray placeholder = ctx.getNDManager().create("");
        placeholder.setName("module_method:generate");
        return new NDList(samples, placeholder);
    }

    /** {@inheritDoc} */
    @Override
    public String processOutput(TranslatorContext ctx, NDList list) throws Exception {
        NDArray result = list.singletonOrThrow();
        StringBuilder sb = new StringBuilder();
//        List<String> sentence = new ArrayList<>();
        for (long ele : result.toLongArray()) {
//            sentence.add(vocabulary.getToken(ele));
            sb.append(vocabulary.getToken(ele));
            if ("<|endoftext|>".equals(vocabulary.getToken(ele))) {
                break;
            }
        }
//        String output = String.join(" ", sentence);
//        return output.replaceAll("[^a-zA-Z0-9<|> ,.!]", "");
        byte[] buf = new byte[sb.length()];
        for (int i = 0; i < sb.length(); ++i) {
            char c = sb.charAt(i);
            buf[i] = BYTES_DECODER.get(c);
        }

        return new String(buf, StandardCharsets.UTF_8);
    }

    /**
     * Returns list of utf-8 byte and a mapping to unicode strings.
     *
     * <p>We specifically avoids mapping to whitespace/control characters the bpe code barfs on. The
     * reversible bpe codes work on unicode strings. This means you need a large # of unicode
     * characters in your vocab if you want to avoid UNKs. When you're at something like a 10B token
     * dataset you end up needing around 5K for decent coverage. This is a significant percentage of
     * your normal, say, 32K bpe vocab. To avoid that, we want lookup tables between utf-8 bytes and
     * unicode strings.
     */
    private static Map<Character, Byte> bpeDecoder() {
        Map<Character, Byte> map = new ConcurrentHashMap<>();
        for (char i = '!'; i <= '~'; ++i) {
            map.put(i, (byte) i);
        }
        for (char i = '¡'; i <= '¬'; ++i) {
            map.put(i, (byte) i);
        }
        for (char i = '®'; i <= 'ÿ'; ++i) {
            map.put(i, (byte) i);
        }

        int n = 0;
        for (char i = 0; i < 256; ++i) {
            if (!map.containsKey(i)) {
                map.put((char) (256 + n), (byte) i);
                ++n;
            }
        }
        return map;
    }

}