vector_database/app.py at main · randomfunction/vector_database · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
fromfastapiimportFastAPI
fromcontextlibimportasynccontextmanager
importfasttext
importcustom_vectordb
importos
importLevenshtein

ft_model=None
db=None

@asynccontextmanager
asyncdeflifespan(app:FastAPI):
globalft_model
globaldb

print("Loading FastText model...")
ft_model=fasttext.load_model("cc.en.300.bin")

# Initialize Engine with strict 300 dimensions implicitly assumed in code structure, mapped to COSINE
db=custom_vectordb.Engine(custom_vectordb.MetricType.COSINE)

ifos.path.exists("snapshot.bin"):
print("Found snapshot.bin! Loading database from disk...")
db.load_from_file("snapshot.bin")
else:
print("snapshot.bin not found. Building database directly from FastText binary model...")
words=ft_model.get_words()
db.reserve(len(words) +1000)
idx=0
forwordinwords:
vec=ft_model.get_word_vector(word).tolist()
db.insert(str(idx),vec,word)
idx+=1
ifidx%1000==0:
print(f"Inserted {idx} words...")
ifidx>=500000:
break
print("Saving snapshot.bin to disk...")
db.save_to_file("snapshot.bin")

print("Database ready!")
yield

fromfastapi.middleware.corsimportCORSMiddleware

app=FastAPI(lifespan=lifespan)

app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
max_age=86400,
)


defsim(v1,v2):
returnLevenshtein.ratio(v1,v2)

@app.get("/search")
defsearch(word:str,k:int, metric:str="cosine"):
search_k=k+10
vec=ft_model.get_word_vector(word).tolist()

results=db.search(vec,search_k)

count=0;

final_results=[]
foriinresults:

if(sim(i.metadata,word)<0.7):
final_results.append(i)
count+=1

if(count==k):
break

results=final_results

foriinfinal_results:
print(i.metadata,sim(i.metadata,word))

return {
"query_vector": [round(v, 3) forvinvec],
"results": [
 {
"word": res.metadata,
"distance": round(float(res.distance), 4),
"vector": [round(v, 3) forvinft_model.get_word_vector(res.metadata).tolist()]

 } forresinresults
 ]
 }