|
from sentence_transformers import SentenceTransformer, util |
|
from PIL import Image |
|
import io |
|
import pickle |
|
import os |
|
import time |
|
import numpy as np |
|
import pprint |
|
|
|
metadata_path = 'conceptarium/metadata.pickle' |
|
|
|
|
|
def init(): |
|
if not os.path.exists(metadata_path): |
|
os.mkdir('conceptarium') |
|
pickle.dump(list(), open(metadata_path, 'wb')) |
|
|
|
|
|
def save(thought): |
|
conceptarium = pickle.load(open(metadata_path, 'rb')) |
|
|
|
if len(conceptarium) > 0: |
|
modality_match = [e.modality == thought.modality for e in conceptarium] |
|
corpus_embeddings = [e.embedding for e in conceptarium] |
|
|
|
results = util.semantic_search( |
|
[thought.embedding], corpus_embeddings, top_k=len(corpus_embeddings), score_function=util.dot_score)[0] |
|
results = [e if modality_match[e['corpus_id']] |
|
else compensate_modality_mismatch(e) for e in results] |
|
|
|
for result in results: |
|
conceptarium[result['corpus_id'] |
|
].interest += result['score'] |
|
|
|
if len(list(filter(lambda x: open(x.filename, 'rb').read() == open(thought.filename, 'rb').read(), conceptarium))) == 0: |
|
conceptarium += [thought] |
|
pickle.dump(conceptarium, open(metadata_path, 'wb')) |
|
|
|
|
|
def find(query, model, relatedness, serendipity, noise, silent, top_k): |
|
conceptarium = pickle.load(open(metadata_path, 'rb')) |
|
|
|
query_embedding = embed(query, model) |
|
query_modality = get_modality(query) |
|
|
|
modality_match = [e.modality == query_modality for e in conceptarium] |
|
corpus_embeddings = [e.embedding for e in conceptarium] |
|
|
|
results = util.semantic_search( |
|
[query_embedding], corpus_embeddings, top_k=len(corpus_embeddings), score_function=util.dot_score)[0] |
|
results = [e if modality_match[e['corpus_id']] |
|
else compensate_modality_mismatch(e) for e in results] |
|
|
|
if not silent: |
|
for result in results: |
|
conceptarium[result['corpus_id'] |
|
].interest += result['score'] |
|
pickle.dump(conceptarium, open(metadata_path, 'wb')) |
|
|
|
for idx, result in enumerate(results): |
|
results[idx]['score'] = (relatedness * result['score'] |
|
- serendipity * |
|
(np.log(conceptarium[result['corpus_id']].interest / (1 - 0.9)) - 0.9 * np.log((time.time() - conceptarium[result['corpus_id']].timestamp) / (3600 * 24) + 0.1))) \ |
|
* np.random.normal(1, noise) |
|
|
|
results = sorted( |
|
results, key=lambda result: result['score'], reverse=True) |
|
memories = [conceptarium[e['corpus_id']] for e in results][:top_k] |
|
return memories |
|
|
|
|
|
def get_doc_paths(directory): |
|
paths = [] |
|
|
|
for root, directories, files in os.walk(directory): |
|
for filename in files: |
|
path = os.path.join(root, filename) |
|
paths.append(path) |
|
|
|
return paths |
|
|
|
|
|
def load_model(): |
|
return SentenceTransformer('clip-ViT-B-32') |
|
|
|
|
|
def embed(content, model): |
|
if get_modality(content) == 'language': |
|
return model.encode(content, convert_to_tensor=True, normalize_embeddings=True) |
|
else: |
|
return model.encode(Image.open(io.BytesIO(content)), convert_to_tensor=True, normalize_embeddings=True) |
|
|
|
|
|
def reset_embeddings(model): |
|
conceptarium = pickle.load(open(metadata_path, 'rb')) |
|
for thought_idx, thought in enumerate(conceptarium): |
|
if thought.modality == 'language': |
|
content = open(thought.filename, 'r').read() |
|
else: |
|
content = open(thought.filename, 'rb').read() |
|
conceptarium[thought_idx].embedding = embed(content, model) |
|
|
|
pickle.dump(conceptarium, open(metadata_path, 'wb')) |
|
|
|
|
|
def get_modality(content): |
|
if isinstance(content, str): |
|
return 'language' |
|
else: |
|
return 'imagery' |
|
|
|
|
|
def compensate_modality_mismatch(result): |
|
result['score'] *= 2.5 |
|
return result |
|
|
|
|
|
class Thought: |
|
def __init__(self, filename, content, model): |
|
self.filename = filename |
|
self.modality = get_modality(content) |
|
self.timestamp = time.time() |
|
self.interest = 1 |
|
self.embedding = embed(content, model) |
|
|
|
def get_content(self): |
|
if self.modality == 'language': |
|
return open(self.filename).read() |
|
elif self.modality == 'imagery': |
|
return open(self.filename, 'rb').read() |
|
|
|
|
|
''' |
|
import json |
|
thoughts = json.load(open('knowledge/base/metadata.json', 'rb')) |
|
|
|
from datetime import datetime |
|
new_thoughts = [] |
|
for thought in thoughts: |
|
new_thought = {} |
|
new_thought['filename'] = thought.filename |
|
new_thought['modality'] = thought.modality |
|
new_thought['timestamp'] = thought.timestamp |
|
new_thought['interest'] = thought.interest |
|
new_thought['embedding'] = thought.embedding |
|
new_thoughts += [new_thought] |
|
|
|
for e_idx, e in enumerate(new_thoughts): |
|
if e['modality'] == 'language': |
|
new_thoughts[e_idx]['modality'] = 'text' |
|
elif e['modality'] == 'imagery': |
|
new_thoughts[e_idx]['modality'] = 'image' |
|
else: |
|
print(e['modality']) |
|
|
|
for e_idx, e in enumerate(new_thoughts): |
|
new_thoughts[e_idx]['embedding'] = e['embedding'].tolist() |
|
|
|
for e_idx, e in enumerate(new_thoughts): |
|
new_thoughts[e_idx]['embedding'] = [round(f, 6) for f in e['embedding']] |
|
|
|
for e_idx, e in enumerate(new_thoughts): |
|
new_thoughts[e_idx]['filename'] = e['filename'].split('/')[-1] |
|
|
|
new_thoughts[0] |
|
|
|
json.dump(new_thoughts, open('knowledge/base/metadata.json', 'wb')) |
|
''' |
|
|
