api.h

Go to the documentation of this file.

#ifndef TINYLLAMA_API_H
#define TINYLLAMA_API_H
 
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>
#include <random>
#include <sstream>
 
#include "model.h"
#include "tokenizer.h"
 
struct KVCache;
 
namespace tinyllama {
 
class TinyLlamaSession {
 public:
  TinyLlamaSession(const std::string& model_path,
                   const std::string& tokenizer_path, int threads = 1,
                   int num_gpu_layers_from_cli = 0, bool cli_use_mmap = true,
                   bool use_kv_quant = false, bool use_batch_generation = false,
                   int max_batch_size = 1);
 
  ~TinyLlamaSession();
 
  std::string generate(const std::string& prompt, int steps = 128,
                      float temperature = 0.1f,
                      int top_k = 40,
                      float top_p = 0.9f,
                      const std::string& system_prompt = "",
                      bool apply_q_a_format = false);
 
  std::vector<std::string> generate_batch(const std::vector<std::string>& prompts,
                                         int steps = 128,
                                         float temperature = 0.1f,
                                         int top_k = 40,
                                         float top_p = 0.9f,
                                         const std::string& system_prompt = "",
                                         bool apply_q_a_format = false);
 
  const Tokenizer* get_tokenizer() const { return tokenizer_.get(); }
  const ModelConfig& get_config() const { return config_; }
  KVCache& get_kv_cache() { return kv_cache_; }
 
 private:
  TinyLlamaSession(const TinyLlamaSession&) = delete;
  TinyLlamaSession& operator=(const TinyLlamaSession&) = delete;
 
  // Parallel batch processing methods
  bool batch_prefill_parallel(const std::vector<std::vector<int>>& all_tokens,
                              const std::vector<int>& prompt_lengths,
                              std::vector<std::vector<float>>& batch_final_logits);
 
  bool batch_generation_parallel(const std::vector<int>& current_tokens,
                                 const std::vector<int>& token_positions,
                                 const std::vector<int>& original_sequence_indices,
                                 std::vector<std::vector<float>>& batch_logits);
 
  std::unique_ptr<TinyLlamaModel> model_;
  std::unique_ptr<Tokenizer> tokenizer_;
  ModelConfig config_;
  KVCache kv_cache_;
  int eos_token_id_;
  std::mt19937 rng_{std::random_device{}()};  // RNG for sampling
  int threads_;
  bool use_batch_generation_; // Enable batch generation
  std::stringstream generated_stream_; // Added for streaming output
  std::string generated_text_for_api_return_; // Added to accumulate full response
  
  // Batch processing support
  int max_batch_size_ = 1;  // Maximum number of sequences for batch processing
};
 
}  // namespace tinyllama
 
#endif