Add sentence transformer for embeddings + config for what models to use

Author: oskarhane

bot.py

@@ -3,7 +3,7 @@
 import streamlit as st
 from langchain.vectorstores.neo4j_vector import Neo4jVector
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.embeddings import OllamaEmbeddings
+from langchain.embeddings import OllamaEmbeddings, SentenceTransformerEmbeddings
 from langchain.chat_models import ChatOpenAI, ChatOllama
 from langchain.chains import RetrievalQAWithSourcesChain
 from langchain.prompts.chat import (
@@ -19,14 +19,30 @@
 username = os.getenv("NEO4J_USERNAME")
 password = os.getenv("NEO4J_PASSWORD")
 ollama_base_url = os.getenv("OLLAMA_BASE_URL")
+embedding_model_name = os.getenv("EMBEDDING_MODEL")
+llm_name = os.getenv("LLM")
 
 os.environ["NEO4J_URL"] = url
 
-# embeddings = OllamaEmbeddings(base_url=ollama_base_url)
-# llm = ChatOllama(temperature=0, base_url=ollama_base_url)
-
-embeddings = OpenAIEmbeddings()
-llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
+if embedding_model_name == "ollama":
+    embeddings = OllamaEmbeddings(base_url=ollama_base_url, model="llama2")
+    print("Embedding: Using Ollama")
+elif embedding_model_name == "openai":
+    embeddings = OpenAIEmbeddings()
+    print("Embedding: Using OpenAI")
+else:
+    embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+    print("Embedding: Using SentenceTransformer")
+    
+if llm_name == "gpt-4":
+    llm = ChatOpenAI(temperature=0, model_name="gpt-4")
+    print("LLM: Using GPT-4")
+elif llm_name == "ollama":
+    llm = ChatOllama(temperature=0, base_url=ollama_base_url, model="llama2")
+    print("LLM: Using Ollama (llama2)")
+else:
+    llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
+    print("LLM: Using GPT-3.5 Turbo")
 
 # LLM only response
 template = "You are a helpful assistant that helps with programming questions."

docker-compose.yml

@@ -28,6 +28,7 @@ services:
       - NEO4J_USERNAME=${NEO4J_USERNAME-neo4j}
       - OPENAI_API_KEY=${OPENAI_API_KEY}
       - OLLAMA_BASE_URL=${OLLAMA_BASE_URL-http://host.docker.internal:11434}
+      - EMBEDDING_MODEL=${EMBEDDING_MODEL-sentence_transformer}
     networks:
       - net
     depends_on:
@@ -53,6 +54,8 @@ services:
       - NEO4J_USERNAME=${NEO4J_USERNAME-neo4j}
       - OPENAI_API_KEY=${OPENAI_API_KEY}
       - OLLAMA_BASE_URL=${OLLAMA_BASE_URL-http://host.docker.internal:11434}
+      - LLM=${LLM-gpt-3.5}
+      - EMBEDDING_MODEL=${EMBEDDING_MODEL-sentence_transformer}
     networks:
       - net
     depends_on:

example.env

@@ -3,3 +3,5 @@ OLLAMA_BASE_URL=http://host.docker.internal:11434
 #NEO4J_URI=neo4j://localhost:7687
 #NEO4J_USERNAME=neo4j
 #NEO4J_PASSWORD=password
+LLM=ollama #or gpt-4 or gpt-3.5
+EMBEDDING_MODEL=sentence_transformer #or openai or ollama

loader.py

@@ -2,7 +2,7 @@
 import requests
 
 from dotenv import load_dotenv
-from langchain.embeddings import OllamaEmbeddings, OpenAIEmbeddings
+from langchain.embeddings import OllamaEmbeddings, OpenAIEmbeddings, SentenceTransformerEmbeddings
 from langchain.graphs import Neo4jGraph
 
 import streamlit as st
@@ -13,18 +13,25 @@
 username = os.getenv("NEO4J_USERNAME")
 password = os.getenv("NEO4J_PASSWORD")
 ollama_base_url = os.getenv("OLLAMA_BASE_URL")
+embedding_model_name = os.getenv("EMBEDDING_MODEL")
 
 os.environ["NEO4J_URL"] = url
 
-# embeddings = OllamaEmbeddings(base_url=ollama_base_url)
-# dimension =  4096 # Ollama
-
-embeddings = OpenAIEmbeddings()
-dimension = 1536  # OpenAi
+if embedding_model_name == "ollama":
+    embeddings = OllamaEmbeddings(base_url=ollama_base_url, model="llama2")
+    dimension =  4096
+    print("Embedding: Using Ollama")
+elif embedding_model_name == "openai":
+    embeddings = OpenAIEmbeddings()
+    dimension = 1536
+    print("Embedding: Using OpenAI")
+else:
+    embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+    dimension = 384
+    print("Embedding: Using SentenceTransformer")
 
 neo4j_graph = Neo4jGraph(url=url, username=username, password=password)
 
-
 def create_constraints():
     neo4j_graph.query(
         "CREATE CONSTRAINT question_id IF NOT EXISTS FOR (q:Question) REQUIRE (q.id) IS UNIQUE"

requirements.txt

@@ -4,4 +4,5 @@ python-dotenv
 wikipedia
 tiktoken
 neo4j
-streamlit
+streamlit
+sentence_transformers