JAX 本身不提供数据加载工具，这使得 JAX 能专注于为构建和训练机器学习模型提供一个高效的计算框架。在 JAX 中加载数据通常借助 TensorFlow 或 PyTorch 来完成。本节将重点介绍如何使用 TensorFlow 在 JAX 中加载数据集。

下面我们开始！

在 JAX 中加载文本数据

我们以“世界杯推文”数据集为例，来演示如何在 JAX 中加载文本数据集。首先导入常用的数据科学库，并查看数据样本。

1import pandas as pd
2
3# 提供解压后文件的路径，读取 CSV 文件
4df = pd.read_csv("/usr/local/notebooks/worldcup-tweets.csv") 
5print(df.head())

清洗文本数据

在使用 TensorFlow 加载数据之前，我们先对数据进行一些预处理。文本处理中的常见步骤是去除停用词。停用词 是指像 a 或 the 这类常见的词汇，它们对判断句子情感倾向帮助不大。NLTK 提供了英文停用词列表，因此我们可以编写函数将其从世界杯推文数据集中移除。

1import nltk
2from nltk.corpus import stopwords
3nltk.download('stopwords', quiet=True)
4
5def remove_stop_words(tweets):
6    tweets_minus_sw = []
7    stop_words = stopwords.words('english')
8    tweets = tweets.split()
9    cleaned_tweets = [tweets_minus_sw.append(word) for word in tweets if word not in stop_words]            
10    cleaned_tweets = ' '.join(tweets_minus_sw)
11    return cleaned_tweets   
12    
13df['tweets'] = df['tweets'].apply(remove_stop_words)
14print(df['tweets'])

标签编码

接下来，我们使用 scikit-learn 中的 LabelEncoder 将情感标签列转换为数值形式。这是因为神经网络只能处理数值型输入。

1from sklearn.preprocessing import LabelEncoder
2
3labelencoder = LabelEncoder()
4df = df.assign(sentiment = labelencoder.fit_transform(df["sentiment"]))
5print(df)

1print("去重前总行数:", len(df))
2
3df = df.drop_duplicates()
4
5print("去重后总行数:", len(df))

1import re
2
3def re_sub(review_sentence):
4    review = re.sub('[^a-zA-Z]', ' ', review_sentence)
5    return review
6
7df['tweets'] = df['tweets'].map(re_sub)
8docs = df['tweets']
9print(docs)