[Generation] Get logits output. (#319)

Author: marvin-Yu

include/models.h

@@ -38,6 +38,8 @@ class Model {
 
     bool isDone();
 
+    std::tuple<float *, int, int> forward();
+
     std::vector<int32_t> generate();
 
     void createSearcher(SearcherConfig &config_);
@@ -50,6 +52,10 @@ class Model {
 
     int getSeqLen() { return seqLen; }
 
+    void setVocabSize(int vocabSize) { this->vocabSize = vocabSize; }
+
+    int getVocabSize() { return this->vocabSize; }
+
     SearcherConfig getConfig() { return configuration; }
 
     void setDecoder(AbstractDecoder *dec);
@@ -70,6 +76,7 @@ class Model {
     std::vector<int32_t> inputIds;
     int batchSize;
     int seqLen;
+    int vocabSize;
     SearcherConfig configuration;
     bool isNewInput;
 };

src/common/transformer_ctx.h

@@ -191,13 +191,6 @@ struct DecoderContext {
     }
 
     void ResetConfigReader(std::string _configPath, std::string _sectionName = "") {
-        fs::path filePath(_configPath);
-
-        if (!fs::exists(filePath)) {
-            printf("Config File %s does not exist!", configPath.c_str());
-            exit(-1);
-        }
-
         this->configPath = _configPath;
         this->configReader = INIReader(_configPath);
         if (this->configReader.ParseError() < 0) {

src/models/models.cpp

@@ -125,6 +125,11 @@ bool Model::isDone() {
     return !isNewInput && searcher->isDone();
 }
 
+std::tuple<float *, int, int> Model::forward() {
+    int64_t dims[3] = {batchSize, 1, seqLen};
+    return decoder->forward(inputIds.data(), dims, 0, true);
+}
+
 std::vector<int32_t> Model::generate() {
     if (inputIds.empty()) {
         printf("Please set input tokens by model.input().\n");
@@ -261,6 +266,7 @@ AutoModel::AutoModel(std::string modelPath, xft::DataType datatype) : Model() {
         exit(-1);
     }
     std::string modeltype = *reader.Sections().begin();
+    setVocabSize(reader.GetInteger(modeltype, "vocab_size"));
 
     if (modeltype == "gpt") {
         switch (datatype) {

src/pytorch/auto_model.h

@@ -129,6 +129,25 @@ struct TorchAutoModel : torch::CustomClassHolder {
                 doSample, temperature, topK, topP, repetitionPenalty, stopWordsList_int32);
     }
 
+    torch::Tensor forward(torch::Tensor &inputIds) {
+        int batchSize = inputIds.size(0);
+        int seqLen = inputIds.size(1);
+        int vocabSize = model->getVocabSize();
+        int logitsN = batchSize * seqLen * vocabSize;
+
+        if (model->getRank() == 0) { input(inputIds); }
+
+        std::tuple<float *, int, int> result = model->forward();
+        float *outBuf = std::get<0>(result);
+        int sampleOffset = std::get<1>(result);
+        int sampleSize = std::get<2>(result);
+
+        // Create a torch::Tensor from the C array
+        int64_t tdims[3] = {batchSize, seqLen, vocabSize};
+        torch::Tensor ret = torch::from_blob(outBuf, tdims, torch::kFloat32);
+        return ret;
+    }
+
     torch::Tensor generate() {
         auto nextTokens = model->generate();
 

src/pytorch/pytorch_warpper.cpp

@@ -22,6 +22,7 @@ TORCH_LIBRARY(xfastertransformer, m) {
             .def("input", &TorchAutoModel::input)
             .def("config", &TorchAutoModel::config)
             .def("is_done", &TorchAutoModel::isDone)
+            .def("forward", &TorchAutoModel::forward)
             .def("generate", &TorchAutoModel::generate)
             .def("finalize", &TorchAutoModel::finalize)
             .def("set_prefix", &TorchAutoModel::setPrefix)

src/utils/weight_util.h

@@ -158,9 +158,10 @@ int loadWeight(std::string filename, T *&ptr, int size, DataType w_type = DataTy
     // By default, read the config.ini configuration file
     // in the same directory as the model file to determine the data type of the file.
     if (w_type == DataType::unknown) {
-        std::filesystem::path pathObj(filename);
-        std::filesystem::path folderPath = pathObj.parent_path();
-        w_type = getWeightType(folderPath.append("config.ini").string());
+        std::size_t pos = filename.find_last_of("/\\");
+        std::string dirPath = filename.substr(0, pos);
+        std::string configFilePath = dirPath + "/config.ini";
+        w_type = getWeightType(configFilePath);
     }
     //1 uint4x2 stores 2 uint4 value, so load size is halfed.
     if constexpr (std::is_same_v<T, uint4x2_t>) { size = size / 2; }

src/xfastertransformer/automodel.py

@@ -38,6 +38,9 @@ def __init__(self, path, dtype: str = "fp16"):
         else:
             raise Exception(f"{self.__class__.__name__} don't support {dtype}.")
 
+    def __call__(self, inputs, **kwargs):
+        return self.model.forward(inputs)
+
     @classmethod
     def from_pretrained(cls, path, dtype: str = "fp16"):
         return cls(path, dtype)