#model_settings.py
import streamlit as st
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from llama_index import LangchainEmbedding, LLMPredictor, PromptHelper, OpenAIEmbedding, ServiceContext 
from llama_index.logger import LlamaLogger
from langchain.chat_models import ChatOpenAI
from langchain import OpenAI
from enum import Enum

class sentenceTransformers(Enum):
    OPTION1 = "sentence-transformers/all-MiniLM-L6-v2"    #default
    OPTION2 = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
    OPTION3 = "sentence-transformers/all-mpnet-base-v2"

def get_sentence_transformer_dropdown():
    options = [e.value for e in sentenceTransformers]
    selected_option = st.selectbox("Sentence transformer:", options)
    return selected_option

def get_embed_model(provider='Langchain', model_name=sentenceTransformers.OPTION1.value):
    # load in HF embedding model from langchain
    embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name=model_name)) if provider=='Langchain' else OpenAIEmbedding()
    return embed_model

def get_prompt_helper():
    # define prompt helper
    max_input_size = 4096
    num_output = 2048
    max_chunk_overlap = 20
    prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)
    return prompt_helper

def get_llm_predictor():
    # define LLM
    num_output = 2048
    
    #llm_predictor = LLMPredictor(llm=OpenAI(temperature=0, model_name="gpt-3.5-turbo", max_tokens=num_output))  
    llm_predictor = LLMPredictor(ChatOpenAI(temperature=0.1, model_name="gpt-3.5-turbo", max_tokens=num_output))
    return llm_predictor

@st.cache_resource
def get_logger():
    llama_logger = LlamaLogger()
    return llama_logger

def get_service_context(llm_predictor=get_llm_predictor(), 
                        embed_model=get_embed_model(), 
                        prompt_helper=get_prompt_helper(), 
                        chunk_size_limit=512, 
                        llama_logger=get_logger()):
    return ServiceContext.from_defaults(llm_predictor=llm_predictor, 
                        embed_model=embed_model, 
                        prompt_helper=prompt_helper, 
                        chunk_size_limit=chunk_size_limit, 
                        llama_logger=llama_logger)

langchain.embeddings.huggingface.HuggingFaceEmbeddings langchain.chat_models.ChatOpenAI
[((705, 751), 'streamlit.selectbox', 'st.selectbox', (['"""Sentence transformer:"""', 'options'], {}), "('Sentence transformer:', options)\n", (717, 751), True, 'import streamlit as st\n'), ((1220, 1279), 'llama_index.PromptHelper', 'PromptHelper', (['max_input_size', 'num_output', 'max_chunk_overlap'], {}), '(max_input_size, num_output, max_chunk_overlap)\n', (1232, 1279), False, 'from llama_index import LangchainEmbedding, LLMPredictor, PromptHelper, OpenAIEmbedding, ServiceContext\n'), ((1684, 1697), 'llama_index.logger.LlamaLogger', 'LlamaLogger', ([], {}), '()\n', (1695, 1697), False, 'from llama_index.logger import LlamaLogger\n'), ((2009, 2192), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'llm_predictor', 'embed_model': 'embed_model', 'prompt_helper': 'prompt_helper', 'chunk_size_limit': 'chunk_size_limit', 'llama_logger': 'llama_logger'}), '(llm_predictor=llm_predictor, embed_model=\n    embed_model, prompt_helper=prompt_helper, chunk_size_limit=\n    chunk_size_limit, llama_logger=llama_logger)\n', (2037, 2192), False, 'from llama_index import LangchainEmbedding, LLMPredictor, PromptHelper, OpenAIEmbedding, ServiceContext\n'), ((1031, 1048), 'llama_index.OpenAIEmbedding', 'OpenAIEmbedding', ([], {}), '()\n', (1046, 1048), False, 'from llama_index import LangchainEmbedding, LLMPredictor, PromptHelper, OpenAIEmbedding, ServiceContext\n'), ((1522, 1600), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0.1)', 'model_name': '"""gpt-3.5-turbo"""', 'max_tokens': 'num_output'}), "(temperature=0.1, model_name='gpt-3.5-turbo', max_tokens=num_output)\n", (1532, 1600), False, 'from langchain.chat_models import ChatOpenAI\n'), ((955, 999), 'langchain.embeddings.huggingface.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {'model_name': 'model_name'}), '(model_name=model_name)\n', (976, 999), False, 'from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n')]