blob: 9afe027e0ebc70ff3beae3913e3bbdbef79c8637 [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 Kupietz4a5e08a2018-06-05 11:07:11 +020047 class Collocator {
48 public:
Marc Kupietzc8ddf452018-01-07 21:33:12 +010049 uint64_t w2;
Marc Kupietz51f93792018-01-25 08:51:01 +010050 uint64_t raw;
Marc Kupietzc8ddf452018-01-07 21:33:12 +010051 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
Marc Kupietzce0b8b02018-06-05 11:06:39 +0200115 // Bouma, Gerlof (2009): <a href="https://svn.spraakdata.gu.se/repos/gerlof/pub/www/Docs/npmi-pfd.pdf">
116 // Normalized (pointwise) mutual information in collocation extraction</a>. In Proceedings of GSCL.
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100117 static inline double ca_npmi(uint64_t f1, uint64_t f2, uint64_t f12, uint64_t total, double window_size) {
Marc Kupietz8caf9912018-06-05 10:51:18 +0200118 if(f12 == 0)
119 return -1.0;
120 else
121 return log2( total * ((double) f12) / (window_size * ((double) f1) * ((double)f2) )) / (-log2(((double) f12 / window_size / total)));
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100122 }
123
124 // Thanopoulos, A., Fakotakis, N., Kokkinakis, G.: Comparative evaluation of collocation extraction metrics.
125 // In: International Conference on Language Resources and Evaluation (LREC-2002). (2002) 620–625
126 // double md = log2(pow((double)max * window_size / total, 2) / (window_size * ((double)_vocab[w1].freq/total) * ((double)_vocab[last_w2].freq/total)));
127 static inline double ca_md(uint64_t f1, uint64_t f2, uint64_t f12, uint64_t total, double window_size) {
128 return log2((double)f12 * f12 / ((double) total * window_size * window_size * f1 * f2));
129 }
130
131 static inline double ca_lfmd(uint64_t f1, uint64_t f2, uint64_t f12, uint64_t total, double window_size) {
Marc Kupietz8caf9912018-06-05 10:51:18 +0200132 if(f12 == 0)
133 return 0;
134 else
135 return log2((double)f12 * f12 / ((double) total * window_size * window_size * f1 * f2)) + log2((double) f12 / window_size / total);
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100136 }
137
Marc Kupietz4b799e92018-01-02 11:04:56 +0100138
139 class CountMergeOperator : public AssociativeMergeOperator {
140 public:
141 CountMergeOperator() {
142 mergeOperator_ = MergeOperators::CreateUInt64AddOperator();
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100143 }
144
Marc Kupietz4b799e92018-01-02 11:04:56 +0100145 virtual bool Merge(const Slice& key,
146 const Slice* existing_value,
147 const Slice& value,
148 std::string* new_value,
149 Logger* logger) const override {
150 assert(new_value->empty());
151 ++num_merge_operator_calls;
152 if (existing_value == nullptr) {
153 new_value->assign(value.data(), value.size());
154 return true;
155 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100156
Marc Kupietz4b799e92018-01-02 11:04:56 +0100157 return mergeOperator_->PartialMerge(
158 key,
159 *existing_value,
160 value,
161 new_value,
162 logger);
163 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100164
Marc Kupietz4b799e92018-01-02 11:04:56 +0100165 virtual const char* Name() const override {
166 return "UInt64AddOperator";
167 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100168
Marc Kupietz4b799e92018-01-02 11:04:56 +0100169 private:
170 std::shared_ptr<MergeOperator> mergeOperator_;
171 };
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100172
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100173
Marc Kupietz4b799e92018-01-02 11:04:56 +0100174 class CollocatorIterator : public Iterator {
175 private:
176 char prefixc[sizeof(uint64_t)];
177 Iterator *base_iterator_;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100178
179
Marc Kupietz4b799e92018-01-02 11:04:56 +0100180 public:
181 CollocatorIterator(Iterator* base_iterator)
182 : base_iterator_(base_iterator)
183 {}
184
Marc Kupietz4b799e92018-01-02 11:04:56 +0100185 void setPrefix(char *prefix) {
186 memcpy(prefixc, prefix, sizeof(uint64_t));
187 }
188
189 virtual void SeekToFirst() { base_iterator_->SeekToFirst(); }
190 virtual void SeekToLast() { base_iterator_->SeekToLast(); }
191 virtual void Seek(const rocksdb::Slice& s) { base_iterator_->Seek(s); }
192 virtual void Prev() { base_iterator_->Prev(); }
193 virtual void Next() { base_iterator_->Next(); }
194 virtual Slice key() const;
195 virtual Slice value() const;
196 virtual Status status() const;
197 virtual bool Valid() const;
198 bool isValid();
199 uint64_t intValue();
200 uint64_t intKey();
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100201
Marc Kupietz4b799e92018-01-02 11:04:56 +0100202 };
Marc Kupietz18375e12017-12-24 10:11:18 +0100203
Marc Kupietz4b799e92018-01-02 11:04:56 +0100204 // rocksdb::CollocatorIterator::CollocatorIterator(Iterator* base_iterator) {}
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100205
Marc Kupietz4b799e92018-01-02 11:04:56 +0100206 bool rocksdb::CollocatorIterator::Valid() const {
Marc Kupietz18375e12017-12-24 10:11:18 +0100207 return base_iterator_->Valid() && key().starts_with(std::string(prefixc,3));
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100208 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100209
Marc Kupietz4b799e92018-01-02 11:04:56 +0100210 bool rocksdb::CollocatorIterator::isValid() {
211 return base_iterator_->Valid() && key().starts_with(std::string(prefixc,3));
Marc Kupietzd31254c2018-01-20 21:29:30 +0100212 // return key().starts_with(std::string(prefixc,3));
Marc Kupietz4b799e92018-01-02 11:04:56 +0100213 }
Marc Kupietz18375e12017-12-24 10:11:18 +0100214
Marc Kupietz4b799e92018-01-02 11:04:56 +0100215 uint64_t rocksdb::CollocatorIterator::intKey() {
216 return DecodeFixed64(base_iterator_->key().data());
217 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100218
Marc Kupietz4b799e92018-01-02 11:04:56 +0100219 uint64_t rocksdb::CollocatorIterator::intValue() {
220 return DecodeFixed64(base_iterator_->value().data());
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100221 }
222
Marc Kupietz37359b12018-01-09 21:11:37 +0100223 class VocabEntry {
224 public:
225 string word;
226 uint64_t freq;
227 };
228
Marc Kupietz6aec7682018-01-10 09:47:48 +0100229 class CollocatorDB {
Marc Kupietz4b799e92018-01-02 11:04:56 +0100230 private:
231 WriteOptions merge_option_; // for merge
232 char _one[sizeof(uint64_t)];
233 Slice _one_slice;
Marc Kupietz37359b12018-01-09 21:11:37 +0100234 vector<VocabEntry> _vocab;
Marc Kupietz4ec51c12019-01-21 11:06:39 +0100235 uint64_t total = 0;
236 uint64_t sentences = 0;
Marc Kupietz37359b12018-01-09 21:11:37 +0100237
Marc Kupietz4b799e92018-01-02 11:04:56 +0100238 protected:
239 std::shared_ptr<DB> db_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100240
Marc Kupietz4b799e92018-01-02 11:04:56 +0100241 WriteOptions put_option_;
242 ReadOptions get_option_;
243 WriteOptions delete_option_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100244
Marc Kupietz4b799e92018-01-02 11:04:56 +0100245 uint64_t default_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100246
Marc Kupietz4b799e92018-01-02 11:04:56 +0100247 std::shared_ptr<DB> OpenDb(const char *dbname);
Marc Kupietz6bb27762018-01-09 17:53:01 +0100248 std::shared_ptr<DB> OpenDbForRead(const char *dbname);
Marc Kupietz37359b12018-01-09 21:11:37 +0100249 void read_vocab(string fname);
250
Marc Kupietz4b799e92018-01-02 11:04:56 +0100251 public:
Marc Kupietz4a5e08a2018-06-05 11:07:11 +0200252 string getWord(uint32_t w1);
Marc Kupietz6aec7682018-01-10 09:47:48 +0100253 CollocatorDB(const char *db_name, bool read_only);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100254
Marc Kupietz6aec7682018-01-10 09:47:48 +0100255 // public interface of CollocatorDB.
Marc Kupietz4b799e92018-01-02 11:04:56 +0100256 // All four functions return false
257 // if the underlying level db operation failed.
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100258
Marc Kupietz4b799e92018-01-02 11:04:56 +0100259 // mapped to a levedb Put
260 bool set(const std::string& key, uint64_t value) {
261 // just treat the internal rep of int64 as the string
262 char buf[sizeof(value)];
263 EncodeFixed64(buf, value);
264 Slice slice(buf, sizeof(value));
265 auto s = db_->Put(put_option_, key, slice);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100266
Marc Kupietz4b799e92018-01-02 11:04:56 +0100267 if (s.ok()) {
268 return true;
269 } else {
270 std::cerr << s.ToString() << std::endl;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100271 return false;
272 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100273 }
Marc Kupietz4b799e92018-01-02 11:04:56 +0100274
275 DB *getDb() {
276 return db_.get();
277 }
278
279 // mapped to a rocksdb Delete
280 bool remove(const std::string& key) {
281 auto s = db_->Delete(delete_option_, key);
282
283 if (s.ok()) {
284 return true;
285 } else {
286 std::cerr << s.ToString() << std::endl;
287 return false;
288 }
289 }
290
291 // mapped to a rocksdb Get
292 bool get(const std::string& key, uint64_t* value) {
293 std::string str;
294 auto s = db_->Get(get_option_, key, &str);
295
296 if (s.IsNotFound()) {
297 // return default value if not found;
298 *value = default_;
299 return true;
300 } else if (s.ok()) {
301 // deserialization
302 if (str.size() != sizeof(uint64_t)) {
303 std::cerr << "value corruption\n";
304 return false;
305 }
306 *value = DecodeFixed64(&str[0]);
307 return true;
308 } else {
309 std::cerr << s.ToString() << std::endl;
310 return false;
311 }
312 }
313
314
315 uint64_t get(const uint32_t w1, const uint32_t w2, const int8_t dist) {
316 char encoded_key[sizeof(uint64_t)];
317 EncodeFixed64(encoded_key, encodeCollocation(w1,w2,dist));
318 uint64_t value = default_;
319 get(std::string(encoded_key, 8), &value);
320 return value;
321 }
322
323 virtual void inc(const std::string& key) {
324 db_->Merge(merge_option_, key, _one_slice);
325 }
326
327 void inc(const uint64_t key) {
328 char encoded_key[sizeof(uint64_t)];
329 EncodeFixed64(encoded_key, key);
330 db_->Merge(merge_option_, std::string(encoded_key, 8), _one_slice);
331 }
332
333 virtual void inc(const uint32_t w1, const uint32_t w2, const uint8_t dist);
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100334 void dump(uint32_t w1, uint32_t w2, int8_t dist);
Marc Kupietz37359b12018-01-09 21:11:37 +0100335 vector<Collocator> get_collocators(uint32_t w1);
Marc Kupietzbd966192018-10-13 14:14:37 +0200336 vector<Collocator> get_collocators(uint32_t w1, uint32_t max_w2);
Marc Kupietz3400aa52018-06-05 10:28:55 +0200337 void dumpSparseLlr(uint32_t w1, uint32_t min_cooccur);
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100338 string collocators2json(vector<Collocator> collocators);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100339
Marc Kupietz4b799e92018-01-02 11:04:56 +0100340 // mapped to a rocksdb Merge operation
341 virtual bool add(const std::string& key, uint64_t value) {
342 char encoded[sizeof(uint64_t)];
343 EncodeFixed64(encoded, value);
344 Slice slice(encoded, sizeof(uint64_t));
345 auto s = db_->Merge(merge_option_, key, slice);
346
347 if (s.ok()) {
348 return true;
349 } else {
350 std::cerr << s.ToString() << std::endl;
351 return false;
352 }
353 }
354
355 CollocatorIterator* SeekIterator(uint64_t w1, uint64_t w2, int8_t dist);
356 };
357
Marc Kupietz6aec7682018-01-10 09:47:48 +0100358 rocksdb::CollocatorDB::CollocatorDB(const char *db_name, bool read_only = false) {
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100359 // merge_option_.sync = true;
Marc Kupietz6bb27762018-01-09 17:53:01 +0100360 if(read_only)
361 db_ = OpenDbForRead(db_name);
362 else
363 db_ = OpenDb(db_name);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100364 assert(db_);
365 uint64_t one = 1;
366 EncodeFixed64(_one, one);
367 _one_slice = Slice(_one, sizeof(uint64_t));
368 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100369
Marc Kupietz6aec7682018-01-10 09:47:48 +0100370 void rocksdb::CollocatorDB::inc(const uint32_t w1, const uint32_t w2, const uint8_t dist) {
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100371 inc(encodeCollocation(w1, w2, dist));
372 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100373
Marc Kupietz6aec7682018-01-10 09:47:48 +0100374 void rocksdb::CollocatorDB::read_vocab(string fname) {
Marc Kupietz37359b12018-01-09 21:11:37 +0100375 char strbuf[2048];
376 uint64_t freq;
377 FILE *fin = fopen(fname.c_str(), "rb");
378 if (fin == NULL) {
379 cout << "Vocabulary file " << fname <<" not found\n";
380 exit(1);
381 }
382 uint64_t i = 0;
383 while(!feof(fin)) {
Marc Kupietzd31254c2018-01-20 21:29:30 +0100384 fscanf(fin, "%s %lu", strbuf, &freq);
Marc Kupietz37359b12018-01-09 21:11:37 +0100385 _vocab.push_back({strbuf, freq});
386 total += freq;
387 i++;
388 }
389 fclose(fin);
Marc Kupietz4ec51c12019-01-21 11:06:39 +0100390
391 char size_fname[256];
392 strcpy(size_fname, fname.c_str());
393 char *pos = strstr(size_fname, ".vocab");
394 if(pos) {
395 *pos=0;
396 strcat(size_fname, ".size");
397 FILE *fp = fopen(size_fname, "r");
398 if (fp != NULL) {
399 fscanf(fp, "%lu", &sentences);
400 fscanf(fp, "%lu", &total);
401 float sl = (float)total/(float)sentences;
402 float w = WINDOW_SIZE;
403 avg_window_size = ((sl > 2*w? (sl-2*w)*2*w: 0) + (double) w * (3*w -1)) / sl;
404 fprintf(stdout, "Size corrections found: corpus size: %lu tokens in %lu sentences, avg. sentence size: %f, avg. window size: %f\n", total, sentences, sl, avg_window_size);
405 fclose(fp);
406 } else {
407 std::cout << "size file " << size_fname << " not found\n";
408 }
409 } else {
410 std::cout << "cannot determine size file " << size_fname << "\n";
411 }
Marc Kupietz37359b12018-01-09 21:11:37 +0100412 }
413
Marc Kupietz6aec7682018-01-10 09:47:48 +0100414 std::shared_ptr<DB> rocksdb::CollocatorDB::OpenDbForRead(const char *name) {
Marc Kupietz6bb27762018-01-09 17:53:01 +0100415 DB* db;
416 Options options;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100417 options.env->SetBackgroundThreads(4);
418 options.create_if_missing = true;
419 options.merge_operator = std::make_shared<CountMergeOperator>();
420 options.max_successive_merges = 0;
421 // options.prefix_extractor.reset(NewFixedPrefixTransform(8));
422 options.IncreaseParallelism();
423 options.OptimizeLevelStyleCompaction();
424 options.prefix_extractor.reset(NewFixedPrefixTransform(3));
Marc Kupietz37359b12018-01-09 21:11:37 +0100425 ostringstream dbname, vocabname;
Marc Kupietz6bb27762018-01-09 17:53:01 +0100426 dbname << name << ".rocksdb";
427 auto s = DB::OpenForReadOnly(options, dbname.str(), &db);
428 if (!s.ok()) {
429 std::cerr << s.ToString() << std::endl;
430 assert(false);
431 }
Marc Kupietz37359b12018-01-09 21:11:37 +0100432 vocabname << name << ".vocab";
433 read_vocab(vocabname.str());
Marc Kupietz6bb27762018-01-09 17:53:01 +0100434 return std::shared_ptr<DB>(db);
435 }
436
Marc Kupietz6aec7682018-01-10 09:47:48 +0100437 std::shared_ptr<DB> rocksdb::CollocatorDB::OpenDb(const char *dbname) {
Marc Kupietz4b799e92018-01-02 11:04:56 +0100438 DB* db;
439 Options options;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100440
441
442 options.env->SetBackgroundThreads(4);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100443 options.create_if_missing = true;
444 options.merge_operator = std::make_shared<CountMergeOperator>();
445 options.max_successive_merges = 0;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100446 // options.prefix_extractor.reset(NewFixedPrefixTransform(8));
447 options.IncreaseParallelism();
448 options.OptimizeLevelStyleCompaction();
449 // options.max_write_buffer_number = 48;
450 // options.max_background_jobs = 48;
451 // options.allow_concurrent_memtable_write=true;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100452 // options.memtable_factory.reset(rocksdb::NewHashLinkListRepFactory(200000));
453 // options.enable_write_thread_adaptive_yield = 1;
454 // options.allow_concurrent_memtable_write = 1;
455 // options.memtable_factory.reset(new rocksdb::SkipListFactory);
456 // options.write_buffer_size = 1 << 22;
457 // options.allow_mmap_reads = true;
458 // options.allow_mmap_writes = true;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100459 // options.max_background_compactions = 40;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100460 // BlockBasedTableOptions table_options;
461 // table_options.filter_policy.reset(NewBloomFilterPolicy(24, false));
462 // options.bloom_locality = 1;
463 // std::shared_ptr<Cache> cache = NewLRUCache(512 * 1024 * 1024);
464 // table_options.block_cache = cache;
465 // options.table_factory.reset(NewBlockBasedTableFactory(table_options));
Marc Kupietz4b799e92018-01-02 11:04:56 +0100466 Status s;
467 // DestroyDB(dbname, Options());
468 s = DB::Open(options, dbname, &db);
469 if (!s.ok()) {
470 std::cerr << s.ToString() << std::endl;
471 assert(false);
472 }
473 return std::shared_ptr<DB>(db);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100474 }
475
Marc Kupietz6aec7682018-01-10 09:47:48 +0100476 CollocatorIterator* rocksdb::CollocatorDB::SeekIterator(uint64_t w1, uint64_t w2, int8_t dist) {
Marc Kupietz18375e12017-12-24 10:11:18 +0100477 ReadOptions options;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100478 options.prefix_same_as_start = true;
Marc Kupietz18375e12017-12-24 10:11:18 +0100479 char prefixc[sizeof(uint64_t)];
480 EncodeFixed64(prefixc, encodeCollocation(w1, w2, dist));
481 Iterator *it = db_->NewIterator(options);
482 CollocatorIterator *cit = new CollocatorIterator(it);
483 cit->Seek(std::string(prefixc,3));// it->Valid() && it->key().starts_with(std::string(prefixc,3)); it->Next()) {
484 cit->setPrefix(prefixc);
485 return cit;
486 }
487
Marc Kupietz6aec7682018-01-10 09:47:48 +0100488 void rocksdb::CollocatorDB::dump(uint32_t w1, uint32_t w2, int8_t dist) {
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100489 auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, w2, dist));
490 for (; it->isValid(); it->Next()) {
491 uint64_t value = it->intValue();
492 uint64_t key = it->intKey();
493 std::cout << "w1:" << W1(key) << ", w2:" << W2(key) << ", dist:" << (int32_t) DIST(key) << " - count:" << value << std::endl;
494 }
495 std::cout << "ready dumping\n";
496 }
497
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100498 double calculateLLR(uint64_t f_X_, uint64_t uintN, uint64_t f_X_Y_, uint64_t f_Y_) {
499 double f_e_, f_o_;
500 double A=0.0, B=0.0, C=0.0, D=0.0, N=0.0;
501 double LLR=0.0, statVal=0.0, minusDiffCoeff=0.0;
502 double BlogB=0.0, ClogC=0.0;
503
504 N = (double)uintN;
505 A = (double)f_X_Y_;
506 B = (double)f_X_ -A;
507 C = (double)f_Y_ -A;
508 D = (double)N -A-B-C;;
509
510 if (B > 0.) BlogB = B*log(B);
511 if (C > 0.) ClogC = C*log(C);
512
513 if ((A>0.) && (D>0.) && (N>0.)) {
514 f_e_ = (double)f_X_ /(double)N;
515 f_o_ = (double)f_X_Y_/(double)f_Y_;
516
517 minusDiffCoeff =
518 ( f_X_==0 ? (double)((-1)*f_X_Y_) :
519 ( f_X_Y_==0 ? (double)((+1)*f_X_) :
520 (f_e_-f_o_)/(f_e_+f_o_)
521 )
522 );
523
524 /* log likelihood ratio */
525 LLR = 2*( A*log(A)
526 +BlogB
527 +ClogC
528 +D*log(D)
529 -(A+B)*log(A+B)
530 -(A+C)*log(A+C)
531 -(B+D)*log(B+D)
532 -(C+D)*log(C+D)
533 +N*log(N)
534 );
535 }
536 return(minusDiffCoeff > 0 ? 0 : (statVal=LLR));
537 }
538
539
540 bool sortByNpmi(const Collocator &lhs, const Collocator &rhs) { return lhs.npmi > rhs.npmi; }
541 bool sortByLfmd(const Collocator &lhs, const Collocator &rhs) { return lhs.lfmd > rhs.lfmd; }
Marc Kupietzd31254c2018-01-20 21:29:30 +0100542 bool sortByLlr(const Collocator &lhs, const Collocator &rhs) { return lhs.llr > rhs.llr; }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100543
Marc Kupietzbd966192018-10-13 14:14:37 +0200544 std::vector<Collocator> rocksdb::CollocatorDB::get_collocators(uint32_t w1, uint32_t max_w2) {
Marc Kupietzd31254c2018-01-20 21:29:30 +0100545 std::vector<Collocator> collocators;
546 uint64_t w2, last_w2 = 0xffffffffffffffff;
Marc Kupietzbd966192018-10-13 14:14:37 +0200547 uint64_t maxv = 0, left = 0, right = 0;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100548 const double window_size = 1;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100549 for ( auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, 0, 0)); it->isValid(); it->Next()) {
550 uint64_t value = it->intValue(),
551 key = it->intKey();
Marc Kupietzbd966192018-10-13 14:14:37 +0200552 if((w2 = W2(key)) > max_w2)
553 continue;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100554 if(last_w2 == 0xffffffffffffffff) last_w2 = w2;
555 if (w2 != last_w2) {
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100556 double pmi = ca_pmi(_vocab[w1].freq, _vocab[last_w2].freq, maxv, total, window_size);
557 double lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, maxv, total, window_size);
558 double left_lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, left, total, 1);
559 double right_lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, right, total, 1);
560 double left_npmi = ca_npmi(_vocab[w1].freq, _vocab[last_w2].freq, left, total, 1);
561 double right_npmi = ca_npmi(_vocab[w1].freq, _vocab[last_w2].freq, right, total, 1);
562 collocators.push_back ( {last_w2, maxv, pmi, pmi / (-log2(((double) maxv / window_size / total))), /* normalize to [-1,1] */
563 calculateLLR(_vocab[w1].freq, total, maxv, _vocab[last_w2].freq), lfmd, pmi*maxv/total/window_size,
564 left_lfmd,
565 right_lfmd,
566 left_npmi,
567 right_npmi}
568 );
Marc Kupietzd31254c2018-01-20 21:29:30 +0100569 last_w2 = w2;
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100570 maxv = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100571 } else {
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100572 if(value > maxv)
573 maxv = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100574 }
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100575 if(DIST(key) == -1)
576 left = value;
577 else if(DIST(key) == 1)
578 right = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100579 }
580
581 sort(collocators.begin(), collocators.end(), sortByLfmd);
582
Marc Kupietz0779a202018-06-05 11:13:35 +0200583 /*
Marc Kupietzd31254c2018-01-20 21:29:30 +0100584 int i=0;
585 for (Collocator c : collocators) {
586 if(i++>10) break;
587 std::cout << "w1:" << _vocab[w1].word << ", w2:" << _vocab[c.w2].word
588 << "\t f(w1):" << _vocab[w1].freq
589 << "\t f(w2):" << _vocab[c.w2].freq
590 << "\t f(w1, x):" << total_w1
Marc Kupietz51f93792018-01-25 08:51:01 +0100591 << "\t f(w1, w2):" << c.raw
Marc Kupietzd31254c2018-01-20 21:29:30 +0100592 << "\t pmi:" << c.pmi
593 << "\t npmi:" << c.npmi
594 << "\t llr:" << c.llr
Marc Kupietzd31254c2018-01-20 21:29:30 +0100595 << "\t lfmd:" << c.lfmd
596 << "\t fpmi:" << c.fpmi
597 << "\t total:" << total
598 << std::endl;
599 }
Marc Kupietz0779a202018-06-05 11:13:35 +0200600 */
Marc Kupietzd31254c2018-01-20 21:29:30 +0100601 return collocators;
602 }
603
Marc Kupietzbd966192018-10-13 14:14:37 +0200604 std::vector<Collocator> rocksdb::CollocatorDB::get_collocators(uint32_t w1) {
605 return get_collocators(w1, UINT32_MAX);
606 }
607
Marc Kupietz3400aa52018-06-05 10:28:55 +0200608 void rocksdb::CollocatorDB::dumpSparseLlr(uint32_t w1, uint32_t min_cooccur) {
609 std::vector<Collocator> collocators;
610 std::stringstream stream;
611 uint64_t w2, last_w2 = 0xffffffffffffffff;
612 uint64_t maxv = 0, total_w1 = 0;
613 bool first = true;
614 for ( auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, 0, 0)); it->isValid(); it->Next()) {
615 uint64_t value = it->intValue(),
616 key = it->intKey();
617 w2 = W2(key);
618 total_w1 += value;
619 if(last_w2 == 0xffffffffffffffff) last_w2 = w2;
620 if (w2 != last_w2) {
621 if(maxv >= min_cooccur) {
622 double llr = calculateLLR(_vocab[w1].freq, total, maxv, _vocab[last_w2].freq);
623 if(first)
624 first = false;
625 else
626 stream << " ";
627 stream << w2 << " " << llr;
628 }
629 last_w2 = w2;
630 maxv = value;
631 } else {
632 if(value > maxv)
633 maxv = value;
634 }
635 }
636 if(first)
637 stream << "1 0.0";
638 stream << "\n";
639 std::cout << stream.str();
640 }
641
Marc Kupietz4b799e92018-01-02 11:04:56 +0100642 rocksdb::Slice rocksdb::CollocatorIterator::key() const { return base_iterator_->key(); }
643 rocksdb::Slice rocksdb::CollocatorIterator::value() const { return base_iterator_->value(); }
644 rocksdb::Status rocksdb::CollocatorIterator::status() const { return base_iterator_->status(); }
645
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100646};
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100647
Marc Kupietz4a5e08a2018-06-05 11:07:11 +0200648string rocksdb::CollocatorDB::getWord(uint32_t w1) {
649 return _vocab[w1].word;
650}
651
Marc Kupietz6aec7682018-01-10 09:47:48 +0100652string rocksdb::CollocatorDB::collocators2json(vector<Collocator> collocators) {
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100653 ostringstream s;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100654 int i = 0;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100655 s << "[";
656 bool first = true;
657 for (Collocator c : collocators) {
Marc Kupietzb999ec52018-06-05 11:20:46 +0200658 if(strncmp(_vocab[c.w2].word.c_str(), "quot", 4) == 0) continue;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100659 if (i++ > 200)
660 break;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100661 if(!first)
662 s << ",\n";
663 else
664 first = false;
665 s << "{"
666 "\"word\":\"" << string(_vocab[c.w2].word) << "\"," <<
667 "\"rank\":" << c.w2 << "," <<
Marc Kupietz51f93792018-01-25 08:51:01 +0100668 "\"f\":" << c.raw << "," <<
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100669 "\"npmi\":" << c.npmi << "," <<
670 "\"llr\":" << c.llr << "," <<
671 "\"lfmd\":" << c.lfmd << "," <<
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100672 "\"fpmi\":" << c.fpmi << "," <<
673 "\"llfmd\":" << c.left_lfmd << "," <<
674 "\"rlfmd\":" << c.right_lfmd << "," <<
675 "\"lnpmi\":" << c.left_npmi << "," <<
676 "\"rnpmi\":" << c.right_npmi <<
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100677 "}";
678 }
679 s << "]\n";
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100680 // cout << s.str();
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100681 return s.str();
682}
683
Marc Kupietz6aec7682018-01-10 09:47:48 +0100684typedef rocksdb::CollocatorDB COLLOCATORS;
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100685
686extern "C" {
Marc Kupietz6aec7682018-01-10 09:47:48 +0100687 COLLOCATORS *open_collocatordb_for_write(char *dbname) {
688 return new rocksdb::CollocatorDB(dbname, false);
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100689 }
690
Marc Kupietz6aec7682018-01-10 09:47:48 +0100691 COLLOCATORS *open_collocatordb(char *dbname) {
692 return new rocksdb::CollocatorDB(dbname, true);
Marc Kupietz6bb27762018-01-09 17:53:01 +0100693 }
694
Marc Kupietz6aec7682018-01-10 09:47:48 +0100695 void inc_collocator(COLLOCATORS *db, uint32_t w1, uint32_t w2, int8_t dist) {
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100696 db->inc(w1, w2, dist);
697 }
698
699 void dump_collocators(COLLOCATORS *db, uint32_t w1, uint32_t w2, int8_t dist) {
700 db->dump(w1, w2, dist);
701 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100702
Marc Kupietz37359b12018-01-09 21:11:37 +0100703 void get_collocators(COLLOCATORS *db, uint32_t w1) {
704 db->get_collocators(w1);
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100705 }
706
Marc Kupietzca3a52e2018-06-05 14:16:23 +0200707 const char *get_word(COLLOCATORS *db, uint32_t w) {
708 return db->getWord(w).c_str();
709 }
710
Marc Kupietz37359b12018-01-09 21:11:37 +0100711 const char *get_collocators_as_json(COLLOCATORS *db, uint32_t w1) {
712 return strdup(db->collocators2json(db->get_collocators(w1)).c_str());
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100713 }
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100714}