update app for tokenizer

app.py

@@ -1,29 +1,30 @@
+import tensorflow as tf
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
+from transformers import AutoTokenizer, TFAutoModel
 
 # Load the dataset
 dataset = load_dataset("sberhe/2023-1000-software-release-notes")
 
-# Explore and preprocess the dataset
-# [Add your data preprocessing steps here]
-
-# Load a pre-trained model and tokenizer
+# Load a pre-trained model and tokenizer (TensorFlow version)
 model_name = "bert-base-uncased"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
+model = TFAutoModel.from_pretrained(model_name)
 
 # Tokenize the data
 def tokenize_function(examples):
-    return tokenizer(examples["text"], padding="max_length", truncation=True)
+    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-# Fine-tuning the model
-training_args = TrainingArguments(output_dir="/results", learning_rate=2e-5, per_device_train_batch_size=16, num_train_epochs=3)
-
-trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_datasets["train"], eval_dataset=tokenized_datasets["test"])
-
-trainer.train()
-
-# Evaluate the model
-# [Add evaluation code here]
+# Function to extract embeddings
+def extract_embeddings(batch):
+    # Convert batch to TensorFlow tensors and correct the dimensions
+    inputs = {k: tf.convert_to_tensor(v) for k, v in batch.items() if k in tokenizer.model_input_names}
+    # Get output from the model
+    outputs = model(**inputs, output_hidden_states=True, return_dict=True)
+    # Extract the embeddings from the last hidden state
+    embeddings = outputs.last_hidden_state
+    return {"embeddings": embeddings.numpy()}
+
+# Apply the function to extract embeddings
+embeddings_dataset = tokenized_datasets.map(extract_embeddings, batched=True)