agent/rag/rag.py

from sentence_transformers import SentenceTransformer, util
from langchain_text_splitters import CharacterTextSplitter
from langchain_core.documents import Document
import pandas as pd


def init_rag(path_file, text_encoder_model):

        # Load Data
        print("loading data...")
        df = pd.read_excel(path_file)  # your Excel file path


        # Preprocess data
        print("Preprocessing data...")
        questions = []
        answers = []

        for _, row in df.iterrows():
            questions.append(row.iloc[0])
            answers.append(row.iloc[1])

        corpus = questions


        # Load Embedder
        print("loading Embedder...")
        model = SentenceTransformer(text_encoder_model)


        corpus_embeddings = model.encode(corpus)

        return model, corpus_embeddings, corpus, answers
    

def get_relevant_question(model, corpus_embeddings, corpus, answers, query:str) -> str:

        # 4) Encode the Arabic query
        query_embedding = model.encode(query)

        # 5) Compute cosine similarity
        cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]

        # 6) Rank results
        top_results = cos_scores.argsort(descending=True)


        matched_result = None
        for idx in top_results:
            if (cos_scores[idx] < .7): # Mostly not relate to her:
                matched_result = None
            else:
                matched_result = (f"Questions {corpus[idx]} \n Answer {answers[idx]} \n (score: {cos_scores[idx]:.4f})")
                print(matched_result)
            break

        return matched_result