blob: eff9c23175d59d89e8ee35e8df4fd852e42abde6 [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 Kupietz8cf7e912019-01-21 17:05:23 +010025#define WINDOW_SIZE 5.0
Marc Kupietz98cbcdc2019-01-21 17:11:27 +010026#define FREQUENCY_THRESHOLD 5
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 Kupietzbbd236e2019-01-21 16:50:19 +0100138 // Evert, Stefan (2004): The Statistics of Word Cooccurrences: Word Pairs and Collocations. PhD dissertation, IMS, University of Stuttgart. Published in 2005, URN urn:nbn:de:bsz:93-opus-23714.
139 // Free PDF available from http://purl.org/stefan.evert/PUB/Evert2004phd.pdf
140 static inline double ca_ll(uint64_t w1, uint64_t w2, uint64_t w12, uint64_t n, uint64_t window_size) {
141 double
142 r1 = (double) w1 * window_size,
143 r2 = (double) n - r1,
144 c1 = w2,
145 c2 = n - c1,
146 o11 = w12, o12 = r1 - o11,
147 o21 = c1 - w12, o22 = r2 - o21,
148 e11 = r1 * c1 / n, e12 = r1 * c2 / n,
149 e21 = r2 * c1 / n, e22 = r2 * c2 / n;
150 return (2 * ( (o11>0? o11 * log(o11/e11):0) + (o12>0? o12 * log(o12/e12):0) + (o21>0? o21 * log(o21/e21):0) + (o22>0? o22 * log(o22/e22):0)));
151 }
Marc Kupietz4b799e92018-01-02 11:04:56 +0100152
153 class CountMergeOperator : public AssociativeMergeOperator {
154 public:
155 CountMergeOperator() {
156 mergeOperator_ = MergeOperators::CreateUInt64AddOperator();
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100157 }
158
Marc Kupietz4b799e92018-01-02 11:04:56 +0100159 virtual bool Merge(const Slice& key,
160 const Slice* existing_value,
161 const Slice& value,
162 std::string* new_value,
163 Logger* logger) const override {
164 assert(new_value->empty());
165 ++num_merge_operator_calls;
166 if (existing_value == nullptr) {
167 new_value->assign(value.data(), value.size());
168 return true;
169 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100170
Marc Kupietz4b799e92018-01-02 11:04:56 +0100171 return mergeOperator_->PartialMerge(
172 key,
173 *existing_value,
174 value,
175 new_value,
176 logger);
177 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100178
Marc Kupietz4b799e92018-01-02 11:04:56 +0100179 virtual const char* Name() const override {
180 return "UInt64AddOperator";
181 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100182
Marc Kupietz4b799e92018-01-02 11:04:56 +0100183 private:
184 std::shared_ptr<MergeOperator> mergeOperator_;
185 };
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100186
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100187
Marc Kupietz4b799e92018-01-02 11:04:56 +0100188 class CollocatorIterator : public Iterator {
189 private:
190 char prefixc[sizeof(uint64_t)];
191 Iterator *base_iterator_;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100192
193
Marc Kupietz4b799e92018-01-02 11:04:56 +0100194 public:
195 CollocatorIterator(Iterator* base_iterator)
196 : base_iterator_(base_iterator)
197 {}
198
Marc Kupietz4b799e92018-01-02 11:04:56 +0100199 void setPrefix(char *prefix) {
200 memcpy(prefixc, prefix, sizeof(uint64_t));
201 }
202
203 virtual void SeekToFirst() { base_iterator_->SeekToFirst(); }
204 virtual void SeekToLast() { base_iterator_->SeekToLast(); }
205 virtual void Seek(const rocksdb::Slice& s) { base_iterator_->Seek(s); }
206 virtual void Prev() { base_iterator_->Prev(); }
207 virtual void Next() { base_iterator_->Next(); }
208 virtual Slice key() const;
209 virtual Slice value() const;
210 virtual Status status() const;
211 virtual bool Valid() const;
212 bool isValid();
213 uint64_t intValue();
214 uint64_t intKey();
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100215
Marc Kupietz4b799e92018-01-02 11:04:56 +0100216 };
Marc Kupietz18375e12017-12-24 10:11:18 +0100217
Marc Kupietz4b799e92018-01-02 11:04:56 +0100218 // rocksdb::CollocatorIterator::CollocatorIterator(Iterator* base_iterator) {}
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100219
Marc Kupietz4b799e92018-01-02 11:04:56 +0100220 bool rocksdb::CollocatorIterator::Valid() const {
Marc Kupietz18375e12017-12-24 10:11:18 +0100221 return base_iterator_->Valid() && key().starts_with(std::string(prefixc,3));
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100222 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100223
Marc Kupietz4b799e92018-01-02 11:04:56 +0100224 bool rocksdb::CollocatorIterator::isValid() {
225 return base_iterator_->Valid() && key().starts_with(std::string(prefixc,3));
Marc Kupietzd31254c2018-01-20 21:29:30 +0100226 // return key().starts_with(std::string(prefixc,3));
Marc Kupietz4b799e92018-01-02 11:04:56 +0100227 }
Marc Kupietz18375e12017-12-24 10:11:18 +0100228
Marc Kupietz4b799e92018-01-02 11:04:56 +0100229 uint64_t rocksdb::CollocatorIterator::intKey() {
230 return DecodeFixed64(base_iterator_->key().data());
231 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100232
Marc Kupietz4b799e92018-01-02 11:04:56 +0100233 uint64_t rocksdb::CollocatorIterator::intValue() {
234 return DecodeFixed64(base_iterator_->value().data());
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100235 }
236
Marc Kupietz37359b12018-01-09 21:11:37 +0100237 class VocabEntry {
238 public:
239 string word;
240 uint64_t freq;
241 };
242
Marc Kupietz6aec7682018-01-10 09:47:48 +0100243 class CollocatorDB {
Marc Kupietz4b799e92018-01-02 11:04:56 +0100244 private:
245 WriteOptions merge_option_; // for merge
246 char _one[sizeof(uint64_t)];
247 Slice _one_slice;
Marc Kupietz37359b12018-01-09 21:11:37 +0100248 vector<VocabEntry> _vocab;
Marc Kupietz4ec51c12019-01-21 11:06:39 +0100249 uint64_t total = 0;
250 uint64_t sentences = 0;
Marc Kupietz8cf7e912019-01-21 17:05:23 +0100251 float avg_window_size = 8.0;
Marc Kupietz37359b12018-01-09 21:11:37 +0100252
Marc Kupietz4b799e92018-01-02 11:04:56 +0100253 protected:
254 std::shared_ptr<DB> db_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100255
Marc Kupietz4b799e92018-01-02 11:04:56 +0100256 WriteOptions put_option_;
257 ReadOptions get_option_;
258 WriteOptions delete_option_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100259
Marc Kupietz4b799e92018-01-02 11:04:56 +0100260 uint64_t default_;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100261
Marc Kupietz4b799e92018-01-02 11:04:56 +0100262 std::shared_ptr<DB> OpenDb(const char *dbname);
Marc Kupietz6bb27762018-01-09 17:53:01 +0100263 std::shared_ptr<DB> OpenDbForRead(const char *dbname);
Marc Kupietz37359b12018-01-09 21:11:37 +0100264 void read_vocab(string fname);
265
Marc Kupietz4b799e92018-01-02 11:04:56 +0100266 public:
Marc Kupietz4a5e08a2018-06-05 11:07:11 +0200267 string getWord(uint32_t w1);
Marc Kupietz6aec7682018-01-10 09:47:48 +0100268 CollocatorDB(const char *db_name, bool read_only);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100269
Marc Kupietz6aec7682018-01-10 09:47:48 +0100270 // public interface of CollocatorDB.
Marc Kupietz4b799e92018-01-02 11:04:56 +0100271 // All four functions return false
272 // if the underlying level db operation failed.
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100273
Marc Kupietz4b799e92018-01-02 11:04:56 +0100274 // mapped to a levedb Put
275 bool set(const std::string& key, uint64_t value) {
276 // just treat the internal rep of int64 as the string
277 char buf[sizeof(value)];
278 EncodeFixed64(buf, value);
279 Slice slice(buf, sizeof(value));
280 auto s = db_->Put(put_option_, key, slice);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100281
Marc Kupietz4b799e92018-01-02 11:04:56 +0100282 if (s.ok()) {
283 return true;
284 } else {
285 std::cerr << s.ToString() << std::endl;
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100286 return false;
287 }
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100288 }
Marc Kupietz4b799e92018-01-02 11:04:56 +0100289
290 DB *getDb() {
291 return db_.get();
292 }
293
294 // mapped to a rocksdb Delete
295 bool remove(const std::string& key) {
296 auto s = db_->Delete(delete_option_, key);
297
298 if (s.ok()) {
299 return true;
300 } else {
301 std::cerr << s.ToString() << std::endl;
302 return false;
303 }
304 }
305
306 // mapped to a rocksdb Get
307 bool get(const std::string& key, uint64_t* value) {
308 std::string str;
309 auto s = db_->Get(get_option_, key, &str);
310
311 if (s.IsNotFound()) {
312 // return default value if not found;
313 *value = default_;
314 return true;
315 } else if (s.ok()) {
316 // deserialization
317 if (str.size() != sizeof(uint64_t)) {
318 std::cerr << "value corruption\n";
319 return false;
320 }
321 *value = DecodeFixed64(&str[0]);
322 return true;
323 } else {
324 std::cerr << s.ToString() << std::endl;
325 return false;
326 }
327 }
328
329
330 uint64_t get(const uint32_t w1, const uint32_t w2, const int8_t dist) {
331 char encoded_key[sizeof(uint64_t)];
332 EncodeFixed64(encoded_key, encodeCollocation(w1,w2,dist));
333 uint64_t value = default_;
334 get(std::string(encoded_key, 8), &value);
335 return value;
336 }
337
338 virtual void inc(const std::string& key) {
339 db_->Merge(merge_option_, key, _one_slice);
340 }
341
342 void inc(const uint64_t key) {
343 char encoded_key[sizeof(uint64_t)];
344 EncodeFixed64(encoded_key, key);
345 db_->Merge(merge_option_, std::string(encoded_key, 8), _one_slice);
346 }
347
348 virtual void inc(const uint32_t w1, const uint32_t w2, const uint8_t dist);
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100349 void dump(uint32_t w1, uint32_t w2, int8_t dist);
Marc Kupietz37359b12018-01-09 21:11:37 +0100350 vector<Collocator> get_collocators(uint32_t w1);
Marc Kupietzbd966192018-10-13 14:14:37 +0200351 vector<Collocator> get_collocators(uint32_t w1, uint32_t max_w2);
Marc Kupietz3400aa52018-06-05 10:28:55 +0200352 void dumpSparseLlr(uint32_t w1, uint32_t min_cooccur);
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100353 string collocators2json(vector<Collocator> collocators);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100354
Marc Kupietz4b799e92018-01-02 11:04:56 +0100355 // mapped to a rocksdb Merge operation
356 virtual bool add(const std::string& key, uint64_t value) {
357 char encoded[sizeof(uint64_t)];
358 EncodeFixed64(encoded, value);
359 Slice slice(encoded, sizeof(uint64_t));
360 auto s = db_->Merge(merge_option_, key, slice);
361
362 if (s.ok()) {
363 return true;
364 } else {
365 std::cerr << s.ToString() << std::endl;
366 return false;
367 }
368 }
369
370 CollocatorIterator* SeekIterator(uint64_t w1, uint64_t w2, int8_t dist);
371 };
372
Marc Kupietz6aec7682018-01-10 09:47:48 +0100373 rocksdb::CollocatorDB::CollocatorDB(const char *db_name, bool read_only = false) {
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100374 // merge_option_.sync = true;
Marc Kupietz6bb27762018-01-09 17:53:01 +0100375 if(read_only)
376 db_ = OpenDbForRead(db_name);
377 else
378 db_ = OpenDb(db_name);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100379 assert(db_);
380 uint64_t one = 1;
381 EncodeFixed64(_one, one);
382 _one_slice = Slice(_one, sizeof(uint64_t));
383 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100384
Marc Kupietz6aec7682018-01-10 09:47:48 +0100385 void rocksdb::CollocatorDB::inc(const uint32_t w1, const uint32_t w2, const uint8_t dist) {
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100386 inc(encodeCollocation(w1, w2, dist));
387 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100388
Marc Kupietz6aec7682018-01-10 09:47:48 +0100389 void rocksdb::CollocatorDB::read_vocab(string fname) {
Marc Kupietz37359b12018-01-09 21:11:37 +0100390 char strbuf[2048];
391 uint64_t freq;
392 FILE *fin = fopen(fname.c_str(), "rb");
393 if (fin == NULL) {
394 cout << "Vocabulary file " << fname <<" not found\n";
395 exit(1);
396 }
397 uint64_t i = 0;
398 while(!feof(fin)) {
Marc Kupietzd31254c2018-01-20 21:29:30 +0100399 fscanf(fin, "%s %lu", strbuf, &freq);
Marc Kupietz37359b12018-01-09 21:11:37 +0100400 _vocab.push_back({strbuf, freq});
401 total += freq;
402 i++;
403 }
404 fclose(fin);
Marc Kupietz4ec51c12019-01-21 11:06:39 +0100405
406 char size_fname[256];
407 strcpy(size_fname, fname.c_str());
408 char *pos = strstr(size_fname, ".vocab");
409 if(pos) {
410 *pos=0;
411 strcat(size_fname, ".size");
412 FILE *fp = fopen(size_fname, "r");
413 if (fp != NULL) {
414 fscanf(fp, "%lu", &sentences);
415 fscanf(fp, "%lu", &total);
416 float sl = (float)total/(float)sentences;
417 float w = WINDOW_SIZE;
418 avg_window_size = ((sl > 2*w? (sl-2*w)*2*w: 0) + (double) w * (3*w -1)) / sl;
419 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);
420 fclose(fp);
421 } else {
422 std::cout << "size file " << size_fname << " not found\n";
423 }
424 } else {
425 std::cout << "cannot determine size file " << size_fname << "\n";
426 }
Marc Kupietz37359b12018-01-09 21:11:37 +0100427 }
428
Marc Kupietz6aec7682018-01-10 09:47:48 +0100429 std::shared_ptr<DB> rocksdb::CollocatorDB::OpenDbForRead(const char *name) {
Marc Kupietz6bb27762018-01-09 17:53:01 +0100430 DB* db;
431 Options options;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100432 options.env->SetBackgroundThreads(4);
433 options.create_if_missing = true;
434 options.merge_operator = std::make_shared<CountMergeOperator>();
435 options.max_successive_merges = 0;
436 // options.prefix_extractor.reset(NewFixedPrefixTransform(8));
437 options.IncreaseParallelism();
438 options.OptimizeLevelStyleCompaction();
439 options.prefix_extractor.reset(NewFixedPrefixTransform(3));
Marc Kupietz37359b12018-01-09 21:11:37 +0100440 ostringstream dbname, vocabname;
Marc Kupietz6bb27762018-01-09 17:53:01 +0100441 dbname << name << ".rocksdb";
442 auto s = DB::OpenForReadOnly(options, dbname.str(), &db);
443 if (!s.ok()) {
444 std::cerr << s.ToString() << std::endl;
445 assert(false);
446 }
Marc Kupietz37359b12018-01-09 21:11:37 +0100447 vocabname << name << ".vocab";
448 read_vocab(vocabname.str());
Marc Kupietz6bb27762018-01-09 17:53:01 +0100449 return std::shared_ptr<DB>(db);
450 }
451
Marc Kupietz6aec7682018-01-10 09:47:48 +0100452 std::shared_ptr<DB> rocksdb::CollocatorDB::OpenDb(const char *dbname) {
Marc Kupietz4b799e92018-01-02 11:04:56 +0100453 DB* db;
454 Options options;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100455
456
457 options.env->SetBackgroundThreads(4);
Marc Kupietz4b799e92018-01-02 11:04:56 +0100458 options.create_if_missing = true;
459 options.merge_operator = std::make_shared<CountMergeOperator>();
460 options.max_successive_merges = 0;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100461 // options.prefix_extractor.reset(NewFixedPrefixTransform(8));
462 options.IncreaseParallelism();
463 options.OptimizeLevelStyleCompaction();
464 // options.max_write_buffer_number = 48;
465 // options.max_background_jobs = 48;
466 // options.allow_concurrent_memtable_write=true;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100467 // options.memtable_factory.reset(rocksdb::NewHashLinkListRepFactory(200000));
468 // options.enable_write_thread_adaptive_yield = 1;
469 // options.allow_concurrent_memtable_write = 1;
470 // options.memtable_factory.reset(new rocksdb::SkipListFactory);
471 // options.write_buffer_size = 1 << 22;
472 // options.allow_mmap_reads = true;
473 // options.allow_mmap_writes = true;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100474 // options.max_background_compactions = 40;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100475 // BlockBasedTableOptions table_options;
476 // table_options.filter_policy.reset(NewBloomFilterPolicy(24, false));
477 // options.bloom_locality = 1;
478 // std::shared_ptr<Cache> cache = NewLRUCache(512 * 1024 * 1024);
479 // table_options.block_cache = cache;
480 // options.table_factory.reset(NewBlockBasedTableFactory(table_options));
Marc Kupietz4b799e92018-01-02 11:04:56 +0100481 Status s;
482 // DestroyDB(dbname, Options());
483 s = DB::Open(options, dbname, &db);
484 if (!s.ok()) {
485 std::cerr << s.ToString() << std::endl;
486 assert(false);
487 }
488 return std::shared_ptr<DB>(db);
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100489 }
490
Marc Kupietz6aec7682018-01-10 09:47:48 +0100491 CollocatorIterator* rocksdb::CollocatorDB::SeekIterator(uint64_t w1, uint64_t w2, int8_t dist) {
Marc Kupietz18375e12017-12-24 10:11:18 +0100492 ReadOptions options;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100493 options.prefix_same_as_start = true;
Marc Kupietz18375e12017-12-24 10:11:18 +0100494 char prefixc[sizeof(uint64_t)];
495 EncodeFixed64(prefixc, encodeCollocation(w1, w2, dist));
496 Iterator *it = db_->NewIterator(options);
497 CollocatorIterator *cit = new CollocatorIterator(it);
498 cit->Seek(std::string(prefixc,3));// it->Valid() && it->key().starts_with(std::string(prefixc,3)); it->Next()) {
499 cit->setPrefix(prefixc);
500 return cit;
501 }
502
Marc Kupietz6aec7682018-01-10 09:47:48 +0100503 void rocksdb::CollocatorDB::dump(uint32_t w1, uint32_t w2, int8_t dist) {
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100504 auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, w2, dist));
505 for (; it->isValid(); it->Next()) {
506 uint64_t value = it->intValue();
507 uint64_t key = it->intKey();
508 std::cout << "w1:" << W1(key) << ", w2:" << W2(key) << ", dist:" << (int32_t) DIST(key) << " - count:" << value << std::endl;
509 }
510 std::cout << "ready dumping\n";
511 }
512
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100513 bool sortByNpmi(const Collocator &lhs, const Collocator &rhs) { return lhs.npmi > rhs.npmi; }
514 bool sortByLfmd(const Collocator &lhs, const Collocator &rhs) { return lhs.lfmd > rhs.lfmd; }
Marc Kupietzd31254c2018-01-20 21:29:30 +0100515 bool sortByLlr(const Collocator &lhs, const Collocator &rhs) { return lhs.llr > rhs.llr; }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100516
Marc Kupietzbd966192018-10-13 14:14:37 +0200517 std::vector<Collocator> rocksdb::CollocatorDB::get_collocators(uint32_t w1, uint32_t max_w2) {
Marc Kupietzd31254c2018-01-20 21:29:30 +0100518 std::vector<Collocator> collocators;
519 uint64_t w2, last_w2 = 0xffffffffffffffff;
Marc Kupietz98cbcdc2019-01-21 17:11:27 +0100520 uint64_t maxv = 0, sum = 0, left = 0, right = 0;
521
Marc Kupietzd31254c2018-01-20 21:29:30 +0100522 for ( auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, 0, 0)); it->isValid(); it->Next()) {
523 uint64_t value = it->intValue(),
524 key = it->intKey();
Marc Kupietzbd966192018-10-13 14:14:37 +0200525 if((w2 = W2(key)) > max_w2)
526 continue;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100527 if(last_w2 == 0xffffffffffffffff) last_w2 = w2;
528 if (w2 != last_w2) {
Marc Kupietz98cbcdc2019-01-21 17:11:27 +0100529 if(sum >= FREQUENCY_THRESHOLD) {
530 double o = sum,
531 r1 = (double)_vocab[w1].freq * avg_window_size,
532 c1 = (double)_vocab[last_w2].freq,
533 e = r1 * c1 / total,
534 pmi = log2(o/e),
535 md = log2(o*o/e),
536 lfmd = log2(o*o*o/e),
537 llr = ca_ll((double)_vocab[w1].freq, (double)_vocab[last_w2].freq, sum, total, avg_window_size);
538 double left_lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, left, total, 1);
539 double right_lfmd = ca_lfmd(_vocab[w1].freq, _vocab[last_w2].freq, right, total, 1);
540 double left_npmi = ca_npmi(_vocab[w1].freq, _vocab[last_w2].freq, left, total, 1);
541 double right_npmi = ca_npmi(_vocab[w1].freq, _vocab[last_w2].freq, right, total, 1);
542 collocators.push_back ( {last_w2, sum, pmi, pmi / (-log2(o)), /* normalize to [-1,1] */
543 llr, lfmd, md,
544 left_lfmd,
545 right_lfmd,
546 left_npmi,
547 right_npmi}
548 );
549 }
Marc Kupietzd31254c2018-01-20 21:29:30 +0100550 last_w2 = w2;
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100551 maxv = value;
Marc Kupietz98cbcdc2019-01-21 17:11:27 +0100552 sum = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100553 } else {
Marc Kupietz98cbcdc2019-01-21 17:11:27 +0100554 sum += value;
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100555 if(value > maxv)
556 maxv = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100557 }
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100558 if(DIST(key) == -1)
559 left = value;
560 else if(DIST(key) == 1)
561 right = value;
Marc Kupietzd31254c2018-01-20 21:29:30 +0100562 }
563
564 sort(collocators.begin(), collocators.end(), sortByLfmd);
565
Marc Kupietz0779a202018-06-05 11:13:35 +0200566 /*
Marc Kupietzd31254c2018-01-20 21:29:30 +0100567 int i=0;
568 for (Collocator c : collocators) {
569 if(i++>10) break;
570 std::cout << "w1:" << _vocab[w1].word << ", w2:" << _vocab[c.w2].word
571 << "\t f(w1):" << _vocab[w1].freq
572 << "\t f(w2):" << _vocab[c.w2].freq
573 << "\t f(w1, x):" << total_w1
Marc Kupietz51f93792018-01-25 08:51:01 +0100574 << "\t f(w1, w2):" << c.raw
Marc Kupietzd31254c2018-01-20 21:29:30 +0100575 << "\t pmi:" << c.pmi
576 << "\t npmi:" << c.npmi
577 << "\t llr:" << c.llr
Marc Kupietzd31254c2018-01-20 21:29:30 +0100578 << "\t lfmd:" << c.lfmd
579 << "\t fpmi:" << c.fpmi
580 << "\t total:" << total
581 << std::endl;
582 }
Marc Kupietz0779a202018-06-05 11:13:35 +0200583 */
Marc Kupietzd31254c2018-01-20 21:29:30 +0100584 return collocators;
585 }
586
Marc Kupietzbd966192018-10-13 14:14:37 +0200587 std::vector<Collocator> rocksdb::CollocatorDB::get_collocators(uint32_t w1) {
588 return get_collocators(w1, UINT32_MAX);
589 }
590
Marc Kupietz3400aa52018-06-05 10:28:55 +0200591 void rocksdb::CollocatorDB::dumpSparseLlr(uint32_t w1, uint32_t min_cooccur) {
592 std::vector<Collocator> collocators;
593 std::stringstream stream;
594 uint64_t w2, last_w2 = 0xffffffffffffffff;
595 uint64_t maxv = 0, total_w1 = 0;
596 bool first = true;
597 for ( auto it = std::unique_ptr<CollocatorIterator>(SeekIterator(w1, 0, 0)); it->isValid(); it->Next()) {
598 uint64_t value = it->intValue(),
599 key = it->intKey();
600 w2 = W2(key);
601 total_w1 += value;
602 if(last_w2 == 0xffffffffffffffff) last_w2 = w2;
603 if (w2 != last_w2) {
604 if(maxv >= min_cooccur) {
Marc Kupietzbbd236e2019-01-21 16:50:19 +0100605 double llr = ca_ll(_vocab[w1].freq, _vocab[last_w2].freq, maxv, total, 1);
Marc Kupietz3400aa52018-06-05 10:28:55 +0200606 if(first)
607 first = false;
608 else
609 stream << " ";
610 stream << w2 << " " << llr;
611 }
612 last_w2 = w2;
613 maxv = value;
614 } else {
615 if(value > maxv)
616 maxv = value;
617 }
618 }
619 if(first)
620 stream << "1 0.0";
621 stream << "\n";
622 std::cout << stream.str();
623 }
624
Marc Kupietz4b799e92018-01-02 11:04:56 +0100625 rocksdb::Slice rocksdb::CollocatorIterator::key() const { return base_iterator_->key(); }
626 rocksdb::Slice rocksdb::CollocatorIterator::value() const { return base_iterator_->value(); }
627 rocksdb::Status rocksdb::CollocatorIterator::status() const { return base_iterator_->status(); }
628
Marc Kupietz28cc53e2017-12-23 17:24:55 +0100629};
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100630
Marc Kupietz4a5e08a2018-06-05 11:07:11 +0200631string rocksdb::CollocatorDB::getWord(uint32_t w1) {
632 return _vocab[w1].word;
633}
634
Marc Kupietz6aec7682018-01-10 09:47:48 +0100635string rocksdb::CollocatorDB::collocators2json(vector<Collocator> collocators) {
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100636 ostringstream s;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100637 int i = 0;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100638 s << "[";
639 bool first = true;
640 for (Collocator c : collocators) {
Marc Kupietzb999ec52018-06-05 11:20:46 +0200641 if(strncmp(_vocab[c.w2].word.c_str(), "quot", 4) == 0) continue;
Marc Kupietz0dd86ef2018-01-11 22:23:17 +0100642 if (i++ > 200)
643 break;
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100644 if(!first)
645 s << ",\n";
646 else
647 first = false;
648 s << "{"
649 "\"word\":\"" << string(_vocab[c.w2].word) << "\"," <<
650 "\"rank\":" << c.w2 << "," <<
Marc Kupietz51f93792018-01-25 08:51:01 +0100651 "\"f\":" << c.raw << "," <<
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100652 "\"npmi\":" << c.npmi << "," <<
653 "\"llr\":" << c.llr << "," <<
654 "\"lfmd\":" << c.lfmd << "," <<
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100655 "\"fpmi\":" << c.fpmi << "," <<
656 "\"llfmd\":" << c.left_lfmd << "," <<
657 "\"rlfmd\":" << c.right_lfmd << "," <<
658 "\"lnpmi\":" << c.left_npmi << "," <<
659 "\"rnpmi\":" << c.right_npmi <<
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100660 "}";
661 }
662 s << "]\n";
Marc Kupietz8e0ebea2018-01-24 09:53:26 +0100663 // cout << s.str();
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100664 return s.str();
665}
666
Marc Kupietz6aec7682018-01-10 09:47:48 +0100667typedef rocksdb::CollocatorDB COLLOCATORS;
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100668
669extern "C" {
Marc Kupietz6aec7682018-01-10 09:47:48 +0100670 COLLOCATORS *open_collocatordb_for_write(char *dbname) {
671 return new rocksdb::CollocatorDB(dbname, false);
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100672 }
673
Marc Kupietz6aec7682018-01-10 09:47:48 +0100674 COLLOCATORS *open_collocatordb(char *dbname) {
675 return new rocksdb::CollocatorDB(dbname, true);
Marc Kupietz6bb27762018-01-09 17:53:01 +0100676 }
677
Marc Kupietz6aec7682018-01-10 09:47:48 +0100678 void inc_collocator(COLLOCATORS *db, uint32_t w1, uint32_t w2, int8_t dist) {
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100679 db->inc(w1, w2, dist);
680 }
681
682 void dump_collocators(COLLOCATORS *db, uint32_t w1, uint32_t w2, int8_t dist) {
683 db->dump(w1, w2, dist);
684 }
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100685
Marc Kupietz37359b12018-01-09 21:11:37 +0100686 void get_collocators(COLLOCATORS *db, uint32_t w1) {
687 db->get_collocators(w1);
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100688 }
689
Marc Kupietzca3a52e2018-06-05 14:16:23 +0200690 const char *get_word(COLLOCATORS *db, uint32_t w) {
691 return db->getWord(w).c_str();
692 }
693
Marc Kupietz37359b12018-01-09 21:11:37 +0100694 const char *get_collocators_as_json(COLLOCATORS *db, uint32_t w1) {
695 return strdup(db->collocators2json(db->get_collocators(w1)).c_str());
Marc Kupietzc8ddf452018-01-07 21:33:12 +0100696 }
Marc Kupietz06c9a9f2018-01-02 16:56:43 +0100697}