File size: 1,196 Bytes
512d794
1d20113
e663f87
 
 
1d20113
 
512d794
e663f87
1d20113
 
512d794
e663f87
1d20113
 
 
512d794
e139833
 
e663f87
 
 
e139833
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# Create FAISS index
def create_faiss_index(texts):
    """
    Create a FAISS index from the provided list of texts.
    """
    import faiss
    from sentence_transformers import SentenceTransformer

    # Load pre-trained SentenceTransformer model
    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    embeddings = model.encode(texts)

    # Create the FAISS index
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)

    return index, texts


# Search the FAISS index
def search_faiss(faiss_index, stored_texts, query, top_k=3):
    """
    Search the FAISS index for the most relevant texts based on the query.
    """
    from sentence_transformers import SentenceTransformer

    # Load the same model used for indexing
    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    # Encode the query into an embedding
    query_embedding = model.encode([query])

    # Search the FAISS index
    distances, indices = faiss_index.search(query_embedding, top_k)

    # Retrieve the corresponding texts
    results = [stored_texts[i] for i in indices[0] if i < len(stored_texts)]

    return results