src/streamlit_app/embedding_model.py

from sentence_transformers import SentenceTransformer, util
from sklearn.linear_model import LogisticRegression
import pickle
from sklearn.model_selection import train_test_split
from sklearn.metrics.pairwise import cosine_similarity
import torch
import pandas as pd
import os

topics = [
    "Ankle Sprain : Recovery exercises for ankle sprains and rehabilitation tips.",
    "Back Pain : Best practices for relieving back pain and strengthening exercises.",
    "Shoulder Injury : Shoulder injury rehabilitation, strengthening, and recovery techniques."
]

def get_encoding_model():
    model_encode = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
    save_model(model_encode,"embedding_model.pkl")

def get_embedding(text):
    if not os.path.exists("embedding_model.pkl"):
        loaded_model = get_encoding_model()

    model_encode = load_model("embedding_model.pkl")
    embedding = model_encode.encode(text)
    return embedding

def save_model(model, filename):
    with open(filename, 'wb') as model_file:
        pickle.dump(model, model_file)
    print(f"Model saved to {filename}")


def load_model(filename):
    # loaded_model = joblib.load('log_reg_model.pkl')
    with open(filename, 'rb') as model_file:
        loaded_model = pickle.load(model_file)
    print(f"Model loaded from {filename}")
    return loaded_model

def get_cosine_similarity(prompt):
    message = "Phsyiotherapy"
    embendding = get_embedding([message,prompt])
    similarity = util.cos_sim(embendding[0], embendding[1]).item()
    return round(similarity, 4)