blob: af9cab9300d068b9ed9a555c9cc507a1c793c34b [file] [log] [blame]
Marc Kupietz28cc53e2017-12-23 17:24:55 +01001#include <typeinfo>
Marc Kupietz4b799e92018-01-02 11:04:56 +01002#define EXPORT __attribute__((visibility("visible")))
3#define IMPORT
Marc Kupietz28cc53e2017-12-23 17:24:55 +01004#include <assert.h>
Marc Kupietz37359b12018-01-09 21:11:37 +01005#include <inttypes.h>
Marc Kupietz28cc53e2017-12-23 17:24:55 +01006#include <memory>
7#include <iostream>
Marc Kupietzc8ddf452018-01-07 21:33:12 +01008#include <algorithm>
9#include <vector>
Marc Kupietz28cc53e2017-12-23 17:24:55 +010010#include <stdint.h>
Marc Kupietzc8ddf452018-01-07 21:33:12 +010011#include <string>
12#include <sstream> // for ostringstream
13#include <math.h>
Marc Kupietzd31254c2018-01-20 21:29:30 +010014#include <rocksdb/cache.h>
Marc Kupietz28cc53e2017-12-23 17:24:55 +010015#include "rocksdb/comparator.h"
16#include "rocksdb/db.h"
17#include "rocksdb/env.h"
Marc Kupietzc8ddf452018-01-07 21:33:12 +010018#include "rocksdb/table.h"
Marc Kupietz28cc53e2017-12-23 17:24:55 +010019#include <rocksdb/merge_operator.h>
Marc Kupietzc8ddf452018-01-07 21:33:12 +010020#include <rocksdb/slice_transform.h>
Marc Kupietz28cc53e2017-12-23 17:24:55 +010021#include "rocksdb/utilities/db_ttl.h"
Marc Kupietzc8ddf452018-01-07 21:33:12 +010022#include "rocksdb/filter_policy.h"
Marc Kupietz28cc53e2017-12-23 17:24:55 +010023#include "merge_operators.h"
24
Marc Kupietzc8ddf452018-01-07 21:33:12 +010025#define AVG_WINDOW_SIZE 7
Marc Kupietz06c9a9f2018-01-02 16:56:43 +010026
Marc Kupietz28cc53e2017-12-23 17:24:55 +010027#define IS_BIG_ENDIAN (*(uint16_t *)"\0\xff" < 0x100)
28#define encodeCollocation(w1, w2, dist) (((uint64_t)dist << 56) | ((uint64_t)w2 << 24) | w1)
Marc Kupietz18375e12017-12-24 10:11:18 +010029#define W1(key) (uint64_t)(key & 0xffffff)
30#define W2(key) (uint64_t)((key >> 24) & 0xffffff)
31#define DIST(key) (int8_t)((uint64_t)((key >> 56) & 0xff))
Marc Kupietzc8ddf452018-01-07 21:33:12 +010032
33typedef struct {
34 uint64_t freq;
35 char *word;
36} vocab_entry;
37
38// typedef struct Collocator {
39// uint64_t w2;
40// uint64_t sum;
41// };
42
Marc Kupietz28cc53e2017-12-23 17:24:55 +010043using namespace rocksdb;
Marc Kupietzc8ddf452018-01-07 21:33:12 +010044using namespace std;
Marc Kupietz28cc53e2017-12-23 17:24:55 +010045
Marc Kupietz4b799e92018-01-02 11:04:56 +010046namespace rocksdb {
Marc Kupietzc8ddf452018-01-07 21:33:12 +010047 class Collocator {
48 public:
49 uint64_t w2;
50 uint64_t sum;
51 double pmi;
52 double npmi;
53 double llr;
Marc Kupietzc8ddf452018-01-07 21:33:12 +010054 double lfmd;
55 double fpmi;
Marc Kupietz8e0ebea2018-01-24 09:53:26 +010056 double left_lfmd;
57 double right_lfmd;
58 double left_npmi;
59 double right_npmi;
Marc Kupietzc8ddf452018-01-07 21:33:12 +010060 };
61
Marc Kupietz28cc53e2017-12-23 17:24:55 +010062 size_t num_merge_operator_calls;
63 void resetNumMergeOperatorCalls() { num_merge_operator_calls = 0; }
Marc Kupietzc8ddf452018-01-07 21:33:12 +010064
Marc Kupietz28cc53e2017-12-23 17:24:55 +010065 size_t num_partial_merge_calls;
66 void resetNumPartialMergeCalls() { num_partial_merge_calls = 0; }
Marc Kupietz28cc53e2017-12-23 17:24:55 +010067
68
Marc Kupietz4b799e92018-01-02 11:04:56 +010069 inline void EncodeFixed64(char* buf, uint64_t value) {
70 if (! IS_BIG_ENDIAN) {
71 memcpy(buf, &value, sizeof(value));
72 } else {
73 buf[0] = value & 0xff;
74 buf[1] = (value >> 8) & 0xff;
75 buf[2] = (value >> 16) & 0xff;
76 buf[3] = (value >> 24) & 0xff;
77 buf[4] = (value >> 32) & 0xff;
78 buf[5] = (value >> 40) & 0xff;
79 buf[6] = (value >> 48) & 0xff;
80 buf[7] = (value >> 56) & 0xff;
81 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +010082 }
83
Marc Kupietz4b799e92018-01-02 11:04:56 +010084 inline uint32_t DecodeFixed32(const char* ptr) {
85 if (! IS_BIG_ENDIAN) {
86 // Load the raw bytes
87 uint32_t result;
88 memcpy(&result, ptr, sizeof(result)); // gcc optimizes this to a plain load
89 return result;
90 } else {
91 return ((static_cast<uint32_t>(static_cast<unsigned char>(ptr[0])))
92 | (static_cast<uint32_t>(static_cast<unsigned char>(ptr[1])) << 8)
93 | (static_cast<uint32_t>(static_cast<unsigned char>(ptr[2])) << 16)
94 | (static_cast<uint32_t>(static_cast<unsigned char>(ptr[3])) << 24));
95 }
96 }
97
98 inline uint64_t DecodeFixed64(const char* ptr) {
99 if (! IS_BIG_ENDIAN) {
100 // Load the raw bytes
101 uint64_t result;
102 memcpy(&result, ptr, sizeof(result)); // gcc optimizes this to a plain load
103 return result;
104 } else {
105 uint64_t lo = DecodeFixed32(ptr);
106 uint64_t hi = DecodeFixed32(ptr + 4);
107 return (hi << 32) | lo;
108 }
109 }
110
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100111 static inline double ca_pmi(uint64_t f1, uint64_t f2, uint64_t f12, uint64_t total, double window_size) {
112 return log2( total * ((double) f12) / (window_size * ((double) f1) * ((double)f2) ));
113 }
114
115 static inline double ca_npmi(uint64_t f1, uint64_t f2, uint64_t f12, uint64_t total, double window_size) {
116 return log2( total * ((double) f12) / (window_size * ((double) f1) * ((double)f2) )) * f12 / total / window_size;
117 }
118
119 // Thanopoulos, A., Fakotakis, N., Kokkinakis, G.: Comparative evaluation of collocation extraction metrics.
120 // In: International Conference on Language Resources and Evaluation (LREC-2002). (2002) 620–625
121 // double md = log2(pow((double)max * window_size / total, 2) / (window_size * ((double)_vocab[w1].freq/total) * ((double)_vocab[last_w2].freq/total)));
122 static inline double ca_md(uint64_t f1, uint64_t f2, uint64_t f12, uint64_t total, double window_size) {
123 return log2((double)f12 * f12 / ((double) total * window_size * window_size * f1 * f2));
124 }
125
126 static inline double ca_lfmd(uint64_t f1, uint64_t f2, uint64_t f12, uint64_t total, double window_size) {
127 return log2((double)f12 * f12 / ((double) total * window_size * window_size * f1 * f2)) + log2((double) f12 / window_size / total);
128 }
129
Marc Kupietz4b799e92018-01-02 11:04:56 +0100130
131 class CountMergeOperator : public AssociativeMergeOperator {
132 public:
133 CountMergeOperator() {
134 mergeOperator_ = MergeOperators::CreateUInt64AddOperator();
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100135 }
136
Marc Kupietz4b799e92018-01-02 11:04:56 +0100137 virtual bool Merge(const Slice& key,
138 const Slice* existing_value,
139 const Slice& value,
140 std::string* new_value,
141 Logger* logger) const override {
142 assert(new_value->empty());
143 ++num_merge_operator_calls;
144 if (existing_value == nullptr) {
145 new_value->assign(value.data(), value.size());
146 return true;
147 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100148
Marc Kupietz4b799e92018-01-02 11:04:56 +0100149 return mergeOperator_->PartialMerge(
150 key,
151 *existing_value,
152 value,
153 new_value,
154 logger);
155 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100156
Marc Kupietz4b799e92018-01-02 11:04:56 +0100157 virtual const char* Name() const override {
158 return "UInt64AddOperator";
159 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100160
Marc Kupietz4b799e92018-01-02 11:04:56 +0100161 private:
162 std::shared_ptr<MergeOperator> mergeOperator_;
163 };
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100164
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100165
Marc Kupietz4b799e92018-01-02 11:04:56 +0100166 class CollocatorIterator : public Iterator {
167 private:
168 char prefixc[sizeof(uint64_t)];
169 Iterator *base_iterator_;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100170
171
Marc Kupietz4b799e92018-01-02 11:04:56 +0100172 public:
173 CollocatorIterator(Iterator* base_iterator)
174 : base_iterator_(base_iterator)
175 {}
176
Marc Kupietz4b799e92018-01-02 11:04:56 +0100177 void setPrefix(char *prefix) {
178 memcpy(prefixc, prefix, sizeof(uint64_t));
179 }
180
181 virtual void SeekToFirst() { base_iterator_->SeekToFirst(); }
182 virtual void SeekToLast() { base_iterator_->SeekToLast(); }
183 virtual void Seek(const rocksdb::Slice& s) { base_iterator_->Seek(s); }
184 virtual void Prev() { base_iterator_->Prev(); }
185 virtual void Next() { base_iterator_->Next(); }
186 virtual Slice key() const;
187 virtual Slice value() const;
188 virtual Status status() const;
189 virtual bool Valid() const;
190 bool isValid();
191 uint64_t intValue();
192 uint64_t intKey();
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100193
Marc Kupietz4b799e92018-01-02 11:04:56 +0100194 };
Marc Kupietz18375e12017-12-24 10:11:18 +0100195
Marc Kupietz4b799e92018-01-02 11:04:56 +0100196 // rocksdb::CollocatorIterator::CollocatorIterator(Iterator* base_iterator) {}
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100197
Marc Kupietz4b799e92018-01-02 11:04:56 +0100198 bool rocksdb::CollocatorIterator::Valid() const {
Marc Kupietz18375e12017-12-24 10:11:18 +0100199 return base_iterator_->Valid() && key().starts_with(std::string(prefixc,3));
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100200 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100201
Marc Kupietz4b799e92018-01-02 11:04:56 +0100202 bool rocksdb::CollocatorIterator::isValid() {
203 return base_iterator_->Valid() && key().starts_with(std::string(prefixc,3));
Marc Kupietzd31254c2018-01-20 21:29:30 +0100204 // return key().starts_with(std::string(prefixc,3));
Marc Kupietz4b799e92018-01-02 11:04:56 +0100205 }
Marc Kupietz18375e12017-12-24 10:11:18 +0100206
Marc Kupietz4b799e92018-01-02 11:04:56 +0100207 uint64_t rocksdb::CollocatorIterator::intKey() {
208 return DecodeFixed64(base_iterator_->key().data());
209 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100210
Marc Kupietz4b799e92018-01-02 11:04:56 +0100211 uint64_t rocksdb::CollocatorIterator::intValue() {
212 return DecodeFixed64(base_iterator_->value().data());
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100213 }
214
Marc Kupietz37359b12018-01-09 21:11:37 +0100215 class VocabEntry {
216 public:
217 string word;
218 uint64_t freq;
219 };
220
Marc Kupietz6aec7682018-01-10 09:47:48 +0100221 class CollocatorDB {
Marc Kupietz4b799e92018-01-02 11:04:56 +0100222 private:
223 WriteOptions merge_option_; // for merge
224 char _one[sizeof(uint64_t)];
225 Slice _one_slice;
Marc Kupietz37359b12018-01-09 21:11:37 +0100226 vector<VocabEntry> _vocab;
227 uint64_t total;
228
Marc Kupietz4b799e92018-01-02 11:04:56 +0100229 protected:
230 std::shared_ptr<DB> db_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100231
Marc Kupietz4b799e92018-01-02 11:04:56 +0100232 WriteOptions put_option_;
233 ReadOptions get_option_;
234 WriteOptions delete_option_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100235
Marc Kupietz4b799e92018-01-02 11:04:56 +0100236 uint64_t default_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100237
Marc Kupietz4b799e92018-01-02 11:04:56 +0100238 std::shared_ptr<DB> OpenDb(const char *dbname);
Marc Kupietz6bb27762018-01-09 17:53:01 +0100239 std::shared_ptr<DB> OpenDbForRead(const char *dbname);
Marc Kupietz37359b12018-01-09 21:11:37 +0100240 void read_vocab(string fname);
241
Marc Kupietz4b799e92018-01-02 11:04:56 +0100242 public:
Marc Kupietz6aec7682018-01-10 09:47:48 +0100243 CollocatorDB(const char *db_name, bool read_only);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100244
Marc Kupietz6aec7682018-01-10 09:47:48 +0100245 // public interface of CollocatorDB.
Marc Kupietz4b799e92018-01-02 11:04:56 +0100246 // All four functions return false
247 // if the underlying level db operation failed.
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100248
Marc Kupietz4b799e92018-01-02 11:04:56 +0100249 // mapped to a levedb Put
250 bool set(const std::string& key, uint64_t value) {
251 // just treat the internal rep of int64 as the string
252 char buf[sizeof(value)];
253 EncodeFixed64(buf, value);
254 Slice slice(buf, sizeof(value));
255 auto s = db_->Put(put_option_, key, slice);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100256
Marc Kupietz4b799e92018-01-02 11:04:56 +0100257 if (s.ok()) {
258 return true;
259 } else {
260 std::cerr << s.ToString() << std::endl;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100261 return false;
262 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100263 }
Marc Kupietz4b799e92018-01-02 11:04:56 +0100264
265 DB *getDb() {
266 return db_.get();
267 }
268
269 // mapped to a rocksdb Delete
270 bool remove(const std::string& key) {
271 auto s = db_->Delete(delete_option_, key);
272
273 if (s.ok()) {
274 return true;
275 } else {
276 std::cerr << s.ToString() << std::endl;
277 return false;
278 }
279 }
280
281 // mapped to a rocksdb Get
282 bool get(const std::string& key, uint64_t* value) {
283 std::string str;
284 auto s = db_->Get(get_option_, key, &str);
285
286 if (s.IsNotFound()) {
287 // return default value if not found;
288 *value = default_;
289 return true;
290 } else if (s.ok()) {
291 // deserialization
292 if (str.size() != sizeof(uint64_t)) {
293 std::cerr << "value corruption\n";
294 return false;
295 }
296 *value = DecodeFixed64(&str[0]);
297 return true;
298 } else {
299 std::cerr << s.ToString() << std::endl;
300 return false;
301 }
302 }
303
304
305 uint64_t get(const uint32_t w1, const uint32_t w2, const int8_t dist) {
306 char encoded_key[sizeof(uint64_t)];
307 EncodeFixed64(encoded_key, encodeCollocation(w1,w2,dist));
308 uint64_t value = default_;
309 get(std::string(encoded_key, 8), &value);
310 return value;
311 }
312
313 virtual void inc(const std::string& key) {
314 db_->Merge(merge_option_, key, _one_slice);
315 }
316
317 void inc(const uint64_t key) {
318 char encoded_key[sizeof(uint64_t)];
319 EncodeFixed64(encoded_key, key);
320 db_->Merge(merge_option_, std::string(encoded_key, 8), _one_slice);
321 }
322
323 virtual void inc(const uint32_t w1, const uint32_t w2, const uint8_t dist);
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100324 void dump(uint32_t w1, uint32_t w2, int8_t dist);
Marc Kupietz37359b12018-01-09 21:11:37 +0100325 vector<Collocator> get_collocators(uint32_t w1);
Marc Kupietzd31254c2018-01-20 21:29:30 +0100326 vector<Collocator> get_collocators_avg(uint32_t w1);
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100327 string collocators2json(vector<Collocator> collocators);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100328
Marc Kupietz4b799e92018-01-02 11:04:56 +0100329 // mapped to a rocksdb Merge operation
330 virtual bool add(const std::string& key, uint64_t value) {
331 char encoded[sizeof(uint64_t)];
332 EncodeFixed64(encoded, value);
333 Slice slice(encoded, sizeof(uint64_t));
334 auto s = db_->Merge(merge_option_, key, slice);
335
336 if (s.ok()) {
337 return true;
338 } else {
339 std::cerr << s.ToString() << std::endl;
340 return false;
341 }
342 }
343
344 CollocatorIterator* SeekIterator(uint64_t w1, uint64_t w2, int8_t dist);
345 };
346
Marc Kupietz6aec7682018-01-10 09:47:48 +0100347 rocksdb::CollocatorDB::CollocatorDB(const char *db_name, bool read_only = false) {
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100348 // merge_option_.sync = true;
Marc Kupietz6bb27762018-01-09 17:53:01 +0100349 if(read_only)
350 db_ = OpenDbForRead(db_name);
351 else
352 db_ = OpenDb(db_name);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100353 assert(db_);
354 uint64_t one = 1;
355 EncodeFixed64(_one, one);
356 _one_slice = Slice(_one, sizeof(uint64_t));
357 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100358
Marc Kupietz6aec7682018-01-10 09:47:48 +0100359 void rocksdb::CollocatorDB::inc(const uint32_t w1, const uint32_t w2, const uint8_t dist) {
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100360 inc(encodeCollocation(w1, w2, dist));
361 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100362
Marc Kupietz6aec7682018-01-10 09:47:48 +0100363 void rocksdb::CollocatorDB::read_vocab(string fname) {
Marc Kupietz37359b12018-01-09 21:11:37 +0100364 char strbuf[2048];
365 uint64_t freq;
366 FILE *fin = fopen(fname.c_str(), "rb");
367 if (fin == NULL) {
368 cout << "Vocabulary file " << fname <<" not found\n";
369 exit(1);
370 }
371 uint64_t i = 0;
372 while(!feof(fin)) {
Marc Kupietzd31254c2018-01-20 21:29:30 +0100373 fscanf(fin, "%s %lu", strbuf, &freq);
Marc Kupietz37359b12018-01-09 21:11:37 +0100374 _vocab.push_back({strbuf, freq});
375 total += freq;
376 i++;
377 }
378 fclose(fin);
379 }
380
Marc Kupietz6aec7682018-01-10 09:47:48 +0100381 std::shared_ptr<DB> rocksdb::CollocatorDB::OpenDbForRead(const char *name) {
Marc Kupietz6bb27762018-01-09 17:53:01 +0100382 DB* db;
383 Options options;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100384 options.env->SetBackgroundThreads(4);
385 options.create_if_missing = true;
386 options.merge_operator = std::make_shared<CountMergeOperator>();
387 options.max_successive_merges = 0;
388 // options.prefix_extractor.reset(NewFixedPrefixTransform(8));
389 options.IncreaseParallelism();
390 options.OptimizeLevelStyleCompaction();
391 options.prefix_extractor.reset(NewFixedPrefixTransform(3));
Marc Kupietz37359b12018-01-09 21:11:37 +0100392 ostringstream dbname, vocabname;
Marc Kupietz6bb27762018-01-09 17:53:01 +0100393 dbname << name << ".rocksdb";
394 auto s = DB::OpenForReadOnly(options, dbname.str(), &db);
395 if (!s.ok()) {
396 std::cerr << s.ToString() << std::endl;
397 assert(false);
398 }
Marc Kupietz37359b12018-01-09 21:11:37 +0100399 vocabname << name << ".vocab";
400 read_vocab(vocabname.str());
Marc Kupietz6bb27762018-01-09 17:53:01 +0100401 return std::shared_ptr<DB>(db);
402 }
403
Marc Kupietz6aec7682018-01-10 09:47:48 +0100404 std::shared_ptr<DB> rocksdb::CollocatorDB::OpenDb(const char *dbname) {
Marc Kupietz4b799e92018-01-02 11:04:56 +0100405 DB* db;
406 Options options;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100407
408
409 options.env->SetBackgroundThreads(4);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100410 options.create_if_missing = true;
411 options.merge_operator = std::make_shared<CountMergeOperator>();
412 options.max_successive_merges = 0;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100413 // options.prefix_extractor.reset(NewFixedPrefixTransform(8));
414 options.IncreaseParallelism();
415 options.OptimizeLevelStyleCompaction();
416 // options.max_write_buffer_number = 48;
417 // options.max_background_jobs = 48;
418 // options.allow_concurrent_memtable_write=true;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100419 // options.memtable_factory.reset(rocksdb::NewHashLinkListRepFactory(200000));
420 // options.enable_write_thread_adaptive_yield = 1;
421 // options.allow_concurrent_memtable_write = 1;
422 // options.memtable_factory.reset(new rocksdb::SkipListFactory);
423 // options.write_buffer_size = 1 << 22;
424 // options.allow_mmap_reads = true;
425 // options.allow_mmap_writes = true;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100426 // options.max_background_compactions = 40;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100427 // BlockBasedTableOptions table_options;
428 // table_options.filter_policy.reset(NewBloomFilterPolicy(24, false));
429 // options.bloom_locality = 1;
430 // std::shared_ptr<Cache> cache = NewLRUCache(512 * 1024 * 1024);
431 // table_options.block_cache = cache;
432 // options.table_factory.reset(NewBlockBasedTableFactory(table_options));
Marc Kupietz4b799e92018-01-02 11:04:56 +0100433 Status s;
434 // DestroyDB(dbname, Options());
435 s = DB::Open(options, dbname, &db);
436 if (!s.ok()) {
437 std::cerr << s.ToString() << std::endl;
438 assert(false);
439 }
440 return std::shared_ptr<DB>(db);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100441 }
442
Marc Kupietz6aec7682018-01-10 09:47:48 +0100443 CollocatorIterator* rocksdb::CollocatorDB::SeekIterator(uint64_t w1, uint64_t w2, int8_t dist) {
Marc Kupietz18375e12017-12-24 10:11:18 +0100444 ReadOptions options;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100445 options.prefix_same_as_start = true;
Marc Kupietz18375e12017-12-24 10:11:18 +0100446 char prefixc[sizeof(uint64_t)];
447 EncodeFixed64(prefixc, encodeCollocation(w1, w2, dist));
448 Iterator *it = db_->NewIterator(options);
449 CollocatorIterator *cit = new CollocatorIterator(it);
450 cit->Seek(std::string(prefixc,3));// it->Valid() && it->key().starts_with(std::string(prefixc,3)); it->Next()) {
451 cit->setPrefix(prefixc);
452 return cit;
453 }
454
Marc Kupietz6aec7682018-01-10 09:47:48 +0100455 void rocksdb::CollocatorDB::dump(uint32_t w1, uint32_t w2, int8_t dist) {
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100456 auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, w2, dist));
457 for (; it->isValid(); it->Next()) {
458 uint64_t value = it->intValue();
459 uint64_t key = it->intKey();
460 std::cout << "w1:" << W1(key) << ", w2:" << W2(key) << ", dist:" << (int32_t) DIST(key) << " - count:" << value << std::endl;
461 }
462 std::cout << "ready dumping\n";
463 }
464
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100465 double calculateLLR(uint64_t f_X_, uint64_t uintN, uint64_t f_X_Y_, uint64_t f_Y_) {
466 double f_e_, f_o_;
467 double A=0.0, B=0.0, C=0.0, D=0.0, N=0.0;
468 double LLR=0.0, statVal=0.0, minusDiffCoeff=0.0;
469 double BlogB=0.0, ClogC=0.0;
470
471 N = (double)uintN;
472 A = (double)f_X_Y_;
473 B = (double)f_X_ -A;
474 C = (double)f_Y_ -A;
475 D = (double)N -A-B-C;;
476
477 if (B > 0.) BlogB = B*log(B);
478 if (C > 0.) ClogC = C*log(C);
479
480 if ((A>0.) && (D>0.) && (N>0.)) {
481 f_e_ = (double)f_X_ /(double)N;
482 f_o_ = (double)f_X_Y_/(double)f_Y_;
483
484 minusDiffCoeff =
485 ( f_X_==0 ? (double)((-1)*f_X_Y_) :
486 ( f_X_Y_==0 ? (double)((+1)*f_X_) :
487 (f_e_-f_o_)/(f_e_+f_o_)
488 )
489 );
490
491 /* log likelihood ratio */
492 LLR = 2*( A*log(A)
493 +BlogB
494 +ClogC
495 +D*log(D)
496 -(A+B)*log(A+B)
497 -(A+C)*log(A+C)
498 -(B+D)*log(B+D)
499 -(C+D)*log(C+D)
500 +N*log(N)
501 );
502 }
503 return(minusDiffCoeff > 0 ? 0 : (statVal=LLR));
504 }
505
506
507 bool sortByNpmi(const Collocator &lhs, const Collocator &rhs) { return lhs.npmi > rhs.npmi; }
508 bool sortByLfmd(const Collocator &lhs, const Collocator &rhs) { return lhs.lfmd > rhs.lfmd; }
Marc Kupietzd31254c2018-01-20 21:29:30 +0100509 bool sortByLlr(const Collocator &lhs, const Collocator &rhs) { return lhs.llr > rhs.llr; }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100510
Marc Kupietzd31254c2018-01-20 21:29:30 +0100511 std::vector<Collocator> rocksdb::CollocatorDB::get_collocators_avg(uint32_t w1) {
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100512 std::vector<Collocator> collocators;
513 uint64_t w2, last_w2 = 0xffffffffffffffff;
514 uint64_t sum = 0, total_w1 = 0;
515 for ( auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, 0, 0)); it->isValid(); it->Next()) {
516 uint64_t value = it->intValue(),
517 key = it->intKey();
518 w2 = W2(key);
519 total_w1 += value;
520 if(last_w2 == 0xffffffffffffffff) last_w2 = w2;
521 if (w2 != last_w2) {
522 double pmi = log2( total * ((double) sum) /
Marc Kupietz37359b12018-01-09 21:11:37 +0100523 (AVG_WINDOW_SIZE * ((double)_vocab[w1].freq) * ((double)_vocab[last_w2].freq) ));
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100524 // Thanopoulos, A., Fakotakis, N., Kokkinakis, G.: Comparative evaluation of collocation extraction metrics. In: International Conference on Language Resources and Evaluation (LREC-2002). (2002) 620–625
Marc Kupietz37359b12018-01-09 21:11:37 +0100525 // double md = log2(pow((double)sum * AVG_WINDOW_SIZE / total, 2) / (AVG_WINDOW_SIZE * ((double)_vocab[w1].freq/total) * ((double)_vocab[last_w2].freq/total)));
526 double md = log2((double)sum * sum / ((double) total * AVG_WINDOW_SIZE * AVG_WINDOW_SIZE * _vocab[w1].freq * _vocab[last_w2].freq));
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100527 collocators.push_back ( {last_w2, sum, pmi, pmi / (-log2(((double) sum / AVG_WINDOW_SIZE / total))), /* normalize to [-1,1] */
Marc Kupietz37359b12018-01-09 21:11:37 +0100528 calculateLLR(_vocab[w1].freq, total, sum, _vocab[last_w2].freq), md, md + log2((double)sum / AVG_WINDOW_SIZE / total), pmi*sum/total/AVG_WINDOW_SIZE} );
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100529 last_w2 = w2;
530 sum = value;
531 } else {
532 sum += value;
533 }
534 }
535
536 sort(collocators.begin(), collocators.end(), sortByLfmd);
537
538 int i=0;
539 for (Collocator c : collocators) {
540 if(i++>10) break;
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100541 std::cout << "dont call me w1:" << _vocab[w1].word << ", w2:" << _vocab[c.w2].word
Marc Kupietz37359b12018-01-09 21:11:37 +0100542 << "\t f(w1):" << _vocab[w1].freq
543 << "\t f(w2):" << _vocab[c.w2].freq
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100544 << "\t f(w1, x):" << total_w1
545 << "\t f(w1, w2):" << c.sum
546 << "\t pmi:" << c.pmi
547 << "\t npmi:" << c.npmi
548 << "\t llr:" << c.llr
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100549 << "\t lfmd:" << c.lfmd
550 << "\t fpmi:" << c.fpmi
551 << "\t total:" << total
552 << std::endl;
553 }
554 return collocators;
555 }
556
Marc Kupietzd31254c2018-01-20 21:29:30 +0100557 std::vector<Collocator> rocksdb::CollocatorDB::get_collocators(uint32_t w1) {
558 std::vector<Collocator> collocators;
559 uint64_t w2, last_w2 = 0xffffffffffffffff;
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100560 uint64_t maxv = 0, left = 0, right = 0, total_w1 = 0;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100561 const double window_size = 1;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100562 for ( auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, 0, 0)); it->isValid(); it->Next()) {
563 uint64_t value = it->intValue(),
564 key = it->intKey();
565 w2 = W2(key);
566 total_w1 += value;
567 if(last_w2 == 0xffffffffffffffff) last_w2 = w2;
568 if (w2 != last_w2) {
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100569 double pmi = ca_pmi(_vocab[w1].freq, _vocab[last_w2].freq, maxv, total, window_size);
570 double lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, maxv, total, window_size);
571 double left_lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, left, total, 1);
572 double right_lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, right, total, 1);
573 double left_npmi = ca_npmi(_vocab[w1].freq, _vocab[last_w2].freq, left, total, 1);
574 double right_npmi = ca_npmi(_vocab[w1].freq, _vocab[last_w2].freq, right, total, 1);
575 collocators.push_back ( {last_w2, maxv, pmi, pmi / (-log2(((double) maxv / window_size / total))), /* normalize to [-1,1] */
576 calculateLLR(_vocab[w1].freq, total, maxv, _vocab[last_w2].freq), lfmd, pmi*maxv/total/window_size,
577 left_lfmd,
578 right_lfmd,
579 left_npmi,
580 right_npmi}
581 );
Marc Kupietzd31254c2018-01-20 21:29:30 +0100582 last_w2 = w2;
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100583 maxv = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100584 } else {
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100585 if(value > maxv)
586 maxv = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100587 }
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100588 if(DIST(key) == -1)
589 left = value;
590 else if(DIST(key) == 1)
591 right = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100592 }
593
594 sort(collocators.begin(), collocators.end(), sortByLfmd);
595
596 int i=0;
597 for (Collocator c : collocators) {
598 if(i++>10) break;
599 std::cout << "w1:" << _vocab[w1].word << ", w2:" << _vocab[c.w2].word
600 << "\t f(w1):" << _vocab[w1].freq
601 << "\t f(w2):" << _vocab[c.w2].freq
602 << "\t f(w1, x):" << total_w1
603 << "\t f(w1, w2):" << c.sum
604 << "\t pmi:" << c.pmi
605 << "\t npmi:" << c.npmi
606 << "\t llr:" << c.llr
Marc Kupietzd31254c2018-01-20 21:29:30 +0100607 << "\t lfmd:" << c.lfmd
608 << "\t fpmi:" << c.fpmi
609 << "\t total:" << total
610 << std::endl;
611 }
612 return collocators;
613 }
614
Marc Kupietz4b799e92018-01-02 11:04:56 +0100615 rocksdb::Slice rocksdb::CollocatorIterator::key() const { return base_iterator_->key(); }
616 rocksdb::Slice rocksdb::CollocatorIterator::value() const { return base_iterator_->value(); }
617 rocksdb::Status rocksdb::CollocatorIterator::status() const { return base_iterator_->status(); }
618
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100619};
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100620
Marc Kupietz6aec7682018-01-10 09:47:48 +0100621string rocksdb::CollocatorDB::collocators2json(vector<Collocator> collocators) {
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100622 ostringstream s;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100623 int i = 0;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100624 s << "[";
625 bool first = true;
626 for (Collocator c : collocators) {
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100627 if (i++ > 200)
628 break;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100629 if(!first)
630 s << ",\n";
631 else
632 first = false;
633 s << "{"
634 "\"word\":\"" << string(_vocab[c.w2].word) << "\"," <<
635 "\"rank\":" << c.w2 << "," <<
636 "\"npmi\":" << c.npmi << "," <<
637 "\"llr\":" << c.llr << "," <<
638 "\"lfmd\":" << c.lfmd << "," <<
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100639 "\"fpmi\":" << c.fpmi << "," <<
640 "\"llfmd\":" << c.left_lfmd << "," <<
641 "\"rlfmd\":" << c.right_lfmd << "," <<
642 "\"lnpmi\":" << c.left_npmi << "," <<
643 "\"rnpmi\":" << c.right_npmi <<
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100644 "}";
645 }
646 s << "]\n";
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100647 // cout << s.str();
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100648 return s.str();
649}
650
Marc Kupietz6aec7682018-01-10 09:47:48 +0100651typedef rocksdb::CollocatorDB COLLOCATORS;
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100652
653extern "C" {
Marc Kupietz6aec7682018-01-10 09:47:48 +0100654 COLLOCATORS *open_collocatordb_for_write(char *dbname) {
655 return new rocksdb::CollocatorDB(dbname, false);
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100656 }
657
Marc Kupietz6aec7682018-01-10 09:47:48 +0100658 COLLOCATORS *open_collocatordb(char *dbname) {
659 return new rocksdb::CollocatorDB(dbname, true);
Marc Kupietz6bb27762018-01-09 17:53:01 +0100660 }
661
Marc Kupietz6aec7682018-01-10 09:47:48 +0100662 void inc_collocator(COLLOCATORS *db, uint32_t w1, uint32_t w2, int8_t dist) {
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100663 db->inc(w1, w2, dist);
664 }
665
666 void dump_collocators(COLLOCATORS *db, uint32_t w1, uint32_t w2, int8_t dist) {
667 db->dump(w1, w2, dist);
668 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100669
Marc Kupietz37359b12018-01-09 21:11:37 +0100670 void get_collocators(COLLOCATORS *db, uint32_t w1) {
671 db->get_collocators(w1);
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100672 }
673
Marc Kupietz37359b12018-01-09 21:11:37 +0100674 const char *get_collocators_as_json(COLLOCATORS *db, uint32_t w1) {
675 return strdup(db->collocators2json(db->get_collocators(w1)).c_str());
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100676 }
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100677}