File size: 668 Bytes
f2ed151
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import gradio as gr
from sentence_transformers import SentenceTransformer, util
import faiss
import numpy as np

model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
memory = []

index = faiss.IndexFlatL2(384)


def add_message(msg):
    emb = model.encode([msg])
    index.add(np.array(emb, dtype='float32'))
    memory.append(msg)
    return f"Memory saved. Total: {len(memory)}"

def recall(query):
    emb = model.encode([query])
    D, I = index.search(np.array(emb, dtype='float32'), 1)
    return memory[I[0][0]] if memory else "Nothing in memory."

ui = gr.Interface(fn=recall, inputs="text", outputs="text", title="Memory Recall")
ui.launch()