1
0
Fork 0
mirror of https://gitlab.com/niansa/libjustlm.git synced 2025-03-06 20:49:17 +01:00

Fixed file type detection

This commit is contained in:
niansa 2023-05-19 17:45:32 +02:00
parent a608135bf7
commit c9dac7cb89
5 changed files with 26 additions and 15 deletions

View file

@ -14,7 +14,9 @@ const LM::Implementation *get_justlm_implementation() {
return &fres; return &fres;
} }
bool magic_match(uint32_t magic) { bool magic_match(std::istream& f) {
uint32_t magic;
f.read(reinterpret_cast<char*>(&magic), sizeof(magic));
return magic == 0x67676d6c; return magic == 0x67676d6c;
} }

View file

@ -9,7 +9,7 @@
static static
Dlhandle get_implementation(uint32_t magic) { Dlhandle get_implementation(std::ifstream& input_f) {
Dlhandle matching; Dlhandle matching;
Dlhandle fallback; Dlhandle fallback;
// Iterate over all libraries // Iterate over all libraries
@ -32,8 +32,9 @@ Dlhandle get_implementation(uint32_t magic) {
continue; continue;
} }
// Set if matching magic // Set if matching magic
auto magic_match = dl.get<bool(uint32_t)>("magic_match"); input_f.seekg(0);
if (magic_match && magic_match(magic)) { auto magic_match = dl.get<bool(std::ifstream&)>("magic_match");
if (magic_match && magic_match(input_f)) {
matching = std::move(dl); matching = std::move(dl);
continue; continue;
} }
@ -48,13 +49,11 @@ LM::Inference *LM::Inference::construct(const std::string &weights_path, const P
static std::vector<Dlhandle> dls; static std::vector<Dlhandle> dls;
// Read magic // Read magic
std::ifstream f(weights_path, std::ios::binary); std::ifstream f(weights_path, std::ios::binary);
uint32_t magic; if (!f) {
if (!f.read(reinterpret_cast<char*>(&magic), sizeof(magic))) {
throw Exception("Failed to open weights file for reading at "+weights_path); throw Exception("Failed to open weights file for reading at "+weights_path);
} }
f.seekg(0);
// Get correct implementation // Get correct implementation
auto impl = get_implementation(magic); auto impl = get_implementation(f);
if (!impl) return nullptr; if (!impl) return nullptr;
// Get inference constructor // Get inference constructor
auto constructor = impl.get<LM::Inference *(const std::string &, std::ifstream&, const LM::Inference::Params &)>("construct"); auto constructor = impl.get<LM::Inference *(const std::string &, std::ifstream&, const LM::Inference::Params &)>("construct");
@ -62,5 +61,6 @@ LM::Inference *LM::Inference::construct(const std::string &weights_path, const P
// Back up Dlhandle // Back up Dlhandle
dls.push_back(std::move(impl)); dls.push_back(std::move(impl));
// Construct inference // Construct inference
f.seekg(0);
return constructor(weights_path, f, p); return constructor(weights_path, f, p);
} }

View file

@ -10,10 +10,21 @@
extern "C" { extern "C" {
const LM::Implementation *get_justlm_implementation() { const LM::Implementation *get_justlm_implementation() {
static LM::Implementation fres{true}; static LM::Implementation fres{false};
return &fres; return &fres;
} }
bool magic_match(std::istream& f) {
// Check magic
uint32_t magic;
f.read(reinterpret_cast<char*>(&magic), sizeof(magic));
if (magic != 0x67676a74) return false;
// Check version
uint32_t version = 0;
f.read(reinterpret_cast<char*>(&version), sizeof(version));
return version >= 2;
}
LM::Inference *construct(const std::string &weights_path, std::ifstream& f, const LM::Inference::Params &p) { LM::Inference *construct(const std::string &weights_path, std::ifstream& f, const LM::Inference::Params &p) {
f.close(); f.close();
return new LM::LLaMAInference(weights_path, p); return new LM::LLaMAInference(weights_path, p);

View file

@ -10,14 +10,10 @@
extern "C" { extern "C" {
const LM::Implementation *get_justlm_implementation() { const LM::Implementation *get_justlm_implementation() {
static LM::Implementation fres{false}; static LM::Implementation fres{true};
return &fres; return &fres;
} }
bool magic_match(uint32_t magic) {
return magic == 0x67676a74;
}
LM::Inference *construct(const std::string &weights_path, std::ifstream& f, const LM::Inference::Params &p) { LM::Inference *construct(const std::string &weights_path, std::ifstream& f, const LM::Inference::Params &p) {
f.close(); f.close();
return new LM::LLaMAInference(weights_path, p); return new LM::LLaMAInference(weights_path, p);

View file

@ -14,7 +14,9 @@ const LM::Implementation *get_justlm_implementation() {
return &fres; return &fres;
} }
bool magic_match(uint32_t magic) { bool magic_match(std::istream& f) {
uint32_t magic;
f.read(reinterpret_cast<char*>(&magic), sizeof(magic));
return magic == 0x67676d6d; return magic == 0x67676d6d;
} }