src/model_setup.py

import os,gc,shutil
from util.conversation_rag import Conversation_RAG
from util.index import *
import torch

class ModelSetup:
    def __init__(self, hf_token, embedding_model, llm):

        self.hf_token = hf_token
        self.embedding_model = embedding_model
        self.llm = llm

    def setup(self):

        if self.embedding_model == "all-roberta-large-v1_1024d":
            embedding_model_repo_id = "sentence-transformers/all-roberta-large-v1"
        elif self.embedding_model == "all-mpnet-base-v2_768d":
            embedding_model_repo_id = "sentence-transformers/all-mpnet-base-v2"


        if self.llm == "Llamav2-7B-Chat":
            llm_repo_id = "meta-llama/Llama-2-7b-chat-hf"
        elif self.llm == "Falcon-7B-Instruct":
            llm_repo_id = "tiiuae/falcon-7b-instruct"


        conv_rag = Conversation_RAG(self.hf_token,
                                    embedding_model_repo_id,
                                    llm_repo_id)

        self.model, self.tokenizer, self.vectordb = conv_rag.load_model_and_tokenizer()
        return "Model Setup Complete"