KNN-CUDAvsCPU/utils.cpp at master · palace22/KNN-CUDAvsCPU · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
#include "utils.h"
#include <assert.h>

namespace pqtable {

FvecsItrReader::FvecsItrReader(std::string filename)
{
    ifs.open(filename, std::ios::binary);
    assert(ifs.is_open());
    Next();
}

bool FvecsItrReader::IsEnd() { return eof_flag; }

std::vector<int> FvecsItrReader::NextInt() { std::vector<int> t(0); return t; }


std::vector<float> FvecsItrReader::Next()
{
    std::vector<float> prev_vec = vec; // return the currently stored vec
    int D;
    if(ifs.read( (char *) &D, sizeof(int) )){ // read "D"
        // Then, read a D-dim vec
        vec.resize(D); // allocate D-dim
        assert(ifs.read( (char *) vec.data(), sizeof(float) * D)); // Read D * float.
        eof_flag = false;
    }else{
        vec.clear();
        eof_flag = true;
    }
    return prev_vec;
}

//////////////////INT//////////
IvecsItrReader::IvecsItrReader(std::string filename)
{
	ifs.open(filename, std::ios::binary);
	assert(ifs.is_open());
	NextInt();
}

bool IvecsItrReader::IsEnd() { return eof_flag; }

std::vector<float> IvecsItrReader::Next() { std::vector<float> t(0); return t; }

std::vector<int> IvecsItrReader::NextInt()
{
	std::vector<int> prev_vec = vec; // return the currently stored vec
	int D;
	if (ifs.read((char *)&D, sizeof(int))) { // read "D"
		// Then, read a D-dim vec
		vec.resize(D); // allocate D-dim
		assert(ifs.read((char *)vec.data(), sizeof(int) * D)); // Read D * int.
		eof_flag = false;
	}
	else {
		vec.clear();
		eof_flag = true;
	}
	return prev_vec;
}

/////////////////

ItrReader::ItrReader(std::string filename, std::string ext){
    if(ext == "fvecs"){
        m_reader = (I_ItrReader *) new FvecsItrReader(filename);
    }else if(ext == "ivecs"){
        m_reader = (I_ItrReader *) new IvecsItrReader(filename);
    }else{
        exit(1);
    }
}

ItrReader::~ItrReader(){
    delete m_reader;
}

bool ItrReader::IsEnd(){
    return m_reader->IsEnd();
}

std::vector<float> ItrReader::Next(){
    return m_reader->Next();
}

std::vector<int> ItrReader::NextInt() {
	return m_reader->NextInt();
}

std::vector<std::vector<float> > ReadTopN(std::string filename, std::string ext, int top_n) {
    std::vector<std::vector<float> > vecs;
    if(top_n != -1){
        vecs.reserve(top_n);
    }
    ItrReader reader(filename, ext);
    while(!reader.IsEnd()){
        if(top_n != -1 && top_n <= (int) vecs.size()){
            break;
        }
        vecs.push_back(reader.Next());
    }
    return vecs;
}

std::vector<std::vector<int> > ReadTopNI(std::string filename, std::string ext, int top_n) {
	std::vector<std::vector<int> > vecs;
	if (top_n != -1) {
		vecs.reserve(top_n);
	}
	ItrReader reader(filename, ext);
	while (!reader.IsEnd()) {
		if (top_n != -1 && top_n <= (int)vecs.size()) {
			break;
		}
		vecs.push_back(reader.NextInt());
	}
	return vecs;
}
}