rag.py

from sentence_transformers import SentenceTransformer, util

# Load model embedding
model = SentenceTransformer('all-MiniLM-L6-v2')

def find_similar_examples(query, df, num_results=5):
    if df.empty:
        print("Empty data!")
        return []

    if "ques" not in df or "ans" not in df:
        print("Column 'ques' or 'ans' not exist in DataFrame.")
        return []

    questions = df["ques"].dropna().tolist()

    if not questions:
        print("There are no valid examples!")
        return []

    question_embeddings = model.encode(questions, convert_to_tensor=True)
    query_embedding = model.encode(query, convert_to_tensor=True)

    scores = util.pytorch_cos_sim(query_embedding, question_embeddings)[0]

    # Limit num of result
    num_results = min(num_results, len(questions))

    top_matches = scores.topk(num_results)

    similar_examples = []
    for idx in top_matches.indices.tolist():  # Convert tensor to list[int]
        question = questions[idx]
        answer = df.iloc[idx]["ans"]
        similar_examples.append({"ques": question, "ans": answer})

    return similar_examples