Switch OAI to tokenizer

GenerateNU · Nov 18, 2023 · 02aa21e · 02aa21e
1 parent aeb979c
commit 02aa21e
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/tokenizer/src/lib.rs b/tokenizer/src/lib.rs
@@ -1,2 +1,2 @@
-pub mod oai;
+pub mod tokenizer;
 pub mod pinecone_utils;
diff --git a/tokenizer/src/main.rs b/tokenizer/src/main.rs
@@ -6,11 +6,11 @@ async fn main() -> std::io::Result<()> {
     mongo.connect().await?;
 
     // set collection to tokenize
-    let collection = "github_data";
-    let oai_key = "generate_2023";
+    let collection: &str = "github_data";
+    let oai_key: &str = "generate_2023";
 
     // TODO: create a new collection for each repo, insert documents into sub collections 
-    let tokenizer: OpenAIClient = OpenAIClient::new(oai_key, mongo, collection);
+    let tokenizer: Tokenizer = Tokenizer::new(oai_key, mongo, collection);
     tokenizer.tokenize_collection(collection);
 
     Ok(())

diff --git a/tokenizer/src/oai.rs → tokenizer/src/tokenizer.rs b/tokenizer/src/oai.rs → tokenizer/src/tokenizer.rs
@@ -5,13 +5,13 @@ use std::error::Error;
 use std::env;
 
 
-pub struct OpenAIClient {
+pub struct Tokenizer {
     oai_client: Client,
     mongo_model: &MongoDriver,
     collection: String,
 }
 
-impl OpenAIClient {
+impl Tokenizer {
 
     pub fn new(openai_api_key: &str, mongo_model: &MongoDriver, collection: &str) -> Self {
         OpenAIClient {
@@ -35,7 +35,7 @@ impl OpenAIClient {
                         let args: openai_rust::embeddings::EmbeddingsArguments = openai_rust::embeddings::EmbeddingsArguments::new("text-embedding-ada-002", text.to_owned());
                         let embedding: Vec<openai_rust::embeddings::EmbeddingsData> = self.oai_client.create_embeddings(args).await.unwrap().data;
 
-                        let update_doc = doc! {
+                        let update_doc: Document = doc! {
                             "$set": { "embedding": embedding }
                         };