Skip to content

Commit

Permalink
support llama3.
Browse files Browse the repository at this point in the history
  • Loading branch information
wangzhaode committed Apr 19, 2024
1 parent 253a41a commit 43e6b00
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 0 deletions.
13 changes: 13 additions & 0 deletions include/llm.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,19 @@ class Yi_6b : public Llama2_7b {
virtual std::vector<int> tokenizer(const std::string& query) override;
virtual bool is_stop(int token_id) override;
};

class Llama3_8b : public Llama2_7b {
public:
Llama3_8b() {
model_name_ = "Llama3_8b";
layer_nums_ = 32;
key_value_shape_ = {2, 1, 8, 0, 128};
hidden_size_ = 4096;
}
private:
virtual std::vector<int> tokenizer(const std::string& query) override;
virtual bool is_stop(int token_id) override;
};
// Llm end

// Embedding start
Expand Down
15 changes: 15 additions & 0 deletions src/llm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,9 @@ Llm* Llm::createLLM(const std::string& path, std::string model_type) {
} else if (model_type.find("yi") != std::string::npos) {
llm = new Yi_6b;
llm->model_name_ = "Yi_6b";
} else if (model_type.find("llama3") != std::string::npos) {
llm = new Llama3_8b;
llm->model_name_ = "Llama3_8b";
}
if (!llm) {
std::cerr << "model type can't judge!" << std::endl;
Expand Down Expand Up @@ -796,6 +799,18 @@ std::vector<int> Yi_6b::tokenizer(const std::string& query) {
bool Yi_6b::is_stop(int token_id) {
return token_id == 7 || token_id == 64001;
}

std::vector<int> Llama3_8b::tokenizer(const std::string& query) {
// <|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n+query+<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n
auto ids = tokenizer_encode(query);
ids.insert(ids.begin(), {128000, 128006, 882, 128007, 271});
ids.insert(ids.end(), {128009, 128006, 78191, 128007, 271});
return ids;
}

bool Llama3_8b::is_stop(int token_id) {
return token_id == 128001 || token_id == 128009;
}
// Llm end

// Embedding start
Expand Down

0 comments on commit 43e6b00

Please sign in to comment.