update embedding model

hello_wordsmith/public_wordsmith_dataset/info.txt

@@ -1,4 +1,114 @@
-Wordsmith is a company based in Scotland. 
-Gigz is CTO.
-Kostis and Derek are ENG. 
-It's pretty lit.
+Wordsmith is a company based in Edinburgh. 
+Derek is the webmaster.
+
+Marketing Info:
+
+Wordsmith
+
+Industries
+
+Join waiting list
+
+
+Secure AI for enterprise legal operations
+Protect your data and harness leading legal AI to make your legal team a delight to interact with across the business. Triage risks, process tasks and expand your capacity.
+
+Join waiting list
+
+Compliance Workflows
+
+Privacy
+
+Commercial
+
+Backed by leading investors
+
+
+Secure Legal AI
+
+Delightfully simple to use
+Meet your team in their workflows and provide instant results
+
+Triage and routing
+
+Triage critical jobs and find the right person
+Analyse thousands of emails, messages and tasks to flag the ones that are critical and need your teams attention.  Mission critical for teams with high volumes of routine interactions with users and their data.
+
+Inbox triage: Set your rules and have your messages ranked by criticality to ensure you get to the right ones first.
+
+Litigation alerts: Scan inbound emails and their attachments to notify you when inbound communications require instant access.
+
+Compliance workflows
+
+Complete compliance workflows in seconds
+Process privacy questionnaires or RFPs, draft DPIAs and respond almost instantly to SARs and privacy queries from users.
+
+DPIA: Accelerate the processing of DPIAs from hours to minutes
+
+SAR: Process subject access requests
+
+Privacy questionnaires: Process repeatable questions rapidly with exceptional quality
+
+Commercial
+
+Accelerate your commerical teams
+Give instant feedback on commercial terms, advanced contract search and trigger contracting workflows
+
+Advanced contract search: Sync thousands of agreements and let anyone in the business get the insight they need
+
+Contract negotiation assistant: Give commercial teams instant feedback on commercial terms
+
+HIGH RISK
+
+MEDIUM RISK
+
+LOW RISK
+
+Unprecedented performance
+
+Engineered for legal
+Wordsmith drives exceptional performance when you are tackling legal issues.  Enabling our enhanced search and query functionality gives your team unique capabilities
+
+How do we benchmark
+
+84%
+
+Wordsmith
+
+63%
+
+Gemini 1.5 Pro
+
+48%
+
+GPT-4 Turbo
+
+45%
+
+Claude 3 Opus
+
+*What percentage of the time these models gave correct answer on our evaluation dataset
+
+Integrations
+
+Wordsmith seamlessly connects with the tools central to your teams workflows
+
+   
+
+Enterprise grade security
+
+Wordsmith is fully SOC2 certified, and our legal agents never train on your data.
+
+Visit trust center
+
+
+AI Consultancy
+
+Making legal legendary
+Speak to our team of AI consultants to help you unlock incredible productivity gains
+
+Join the waiting list
+
+Wordsmith
+
+AI powered legal operations for enterprise

hello_wordsmith/wordsmith.py

@@ -3,18 +3,29 @@ import sys
 
 import chromadb
 from llama_index.cli.rag import RagCLI
-from llama_index.core import (ChatPromptTemplate, SimpleDirectoryReader,
-                              StorageContext, VectorStoreIndex)
+from llama_index.core import (
+    ChatPromptTemplate,
+    SimpleDirectoryReader,
+    VectorStoreIndex
+)
 from llama_index.core.base.llms.types import ChatMessage, MessageRole
 from llama_index.core.ingestion import IngestionPipeline
-from llama_index.core.prompts.base import ChatPromptTemplate
 from llama_index.core.query_pipeline import InputComponent, QueryPipeline
 from llama_index.core.response_synthesizers import TreeSummarize
-from llama_index.core.storage.docstore import SimpleDocumentStore
+from llama_index.core import Settings
+from llama_index.embeddings.openai import (
+    OpenAIEmbedding,
+    OpenAIEmbeddingModelType
+)
 from llama_index.llms.openai import OpenAI
 from llama_index.vector_stores.chroma import ChromaVectorStore
 
 
+Settings.embed_model = OpenAIEmbedding(
+    model="text-embedding-3-small"
+)
+
+
 def initialize_chroma_db():
     chroma_client = chromadb.EphemeralClient()
     chroma_collection = chroma_client.create_collection("wordsmith")
@@ -27,8 +38,7 @@ def setup_document_storage(vector_store):
     dataset_path = os.path.join(package_directory, "public_wordsmith_dataset")
     reader = SimpleDirectoryReader(input_dir=dataset_path)
     docs = reader.load_data()
-    storage_context = StorageContext.from_defaults(vector_store=vector_store)
-    index = VectorStoreIndex.from_documents(docs, storage_context=storage_context)
+    index = VectorStoreIndex.from_documents(docs)
     return index
 
 
@@ -79,7 +89,7 @@ def configure_query_pipeline(index, llm):
     )
     query_pipeline = QueryPipeline()
 
-    retriever = index.as_retriever(similarity_top_k=5)
+    retriever = index.as_retriever(similarity_top_k=20)
     summarizer = TreeSummarize(
         llm=llm, streaming=True, summary_template=text_qa_chat_template
     )