FeatureExtraction_Train_Test_Split.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04cd7c3d-7c37-482d-81ba-31dd6e4ec666",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import random\n",
    "import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import MultipleLocator\n",
    "import matplotlib.dates as mdates\n",
    "import seaborn as sns\n",
    "import math\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "from collections import defaultdict\n",
    "from sklearn.metrics import mean_absolute_error, mean_squared_error\n",
    "from abc import ABC, abstractmethod\n",
    "import re\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "import pickle\n",
    "import time\n",
    "import itertools\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88c73c5b-e695-4731-aa60-9f2f230a9194",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "from pandas.core.common import SettingWithCopyWarning\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=RuntimeWarning) \n",
    "warnings.simplefilter(action=\"ignore\", category=SettingWithCopyWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d971941f-6d55-4de6-8643-4fa4c852ef07",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e3b36e4-d3e6-4e4c-b171-5f81a9c065bb",
   "metadata": {},
   "source": [
    "# Feature Engineering and Train Test Split:\n",
    "\n",
    "* This notebook contains code that was utilized for preprocessing user timelines to extract features for our forecasting models\n",
    "* It also has code related to how the train,val,test split was conducted"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b9051e8-823a-45a8-a4ed-325f82f79aec",
   "metadata": {},
   "source": [
    "## Extractors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fdecd57-93c9-4a44-9e3c-378ed294b7c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeatureExtractor(ABC):\n",
    "    \n",
    "    @abstractmethod\n",
    "    def extract(self,bin_ts,**kwargs):\n",
    "        \"\"\"\n",
    "        Extracts Features for each user sequence\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "class EngagementCountExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.ps_index_map = {k:i for i,k in enumerate(range(-3,4,1))}\n",
    "        \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        Extract Count Vector of engagements for a given time step bin\n",
    "        \n",
    "        * Default : [1,1,1,1,1,1,1] (represents no engagement)\n",
    "        * Shifted by 1 for evaulation metrics to work properly (smape and nmse)\n",
    "        \"\"\"\n",
    "        df = bin_ts[\"eng_dataframe\"]\n",
    "        ps_vec = [1.0]*7\n",
    "        if len(df)>0:\n",
    "            mps = df[\"matched_partisans\"].tolist()\n",
    "            mps = flatten(mps)\n",
    "            for p in mps:\n",
    "                ps_vec[self.ps_index_map[p]]+=1\n",
    "        \n",
    "        return ps_vec\n",
    "\n",
    "class EngagementTypeCountExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "    \n",
    "    def get_engagement_type(self,urow):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        if len(urow.matched_urls)>0 and len(urow.matched_mentions)<=0:\n",
    "            return 0\n",
    "        elif len(urow.matched_mentions)>0 and len(urow.matched_urls)<=0:\n",
    "            return 1\n",
    "        elif len(urow.matched_urls)>0 and len(urow.matched_mentions)>0:\n",
    "            return 2\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        df = bin_ts[\"eng_dataframe\"]\n",
    "        engtype_vec = [0.0]*3\n",
    "        if len(df)>0:\n",
    "            df[\"eng_type\"] = df.apply(lambda x: self.get_engagement_type(x),axis=1)\n",
    "            eng_types = df[\"eng_type\"].tolist()\n",
    "            for e in eng_types:\n",
    "                engtype_vec[e]+=1\n",
    "        \n",
    "        return engtype_vec\n",
    "        \n",
    "        \n",
    "    \n",
    "class EngagementPublicMetricsExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        pass\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        public metrics types:\n",
    "        * retweet count\n",
    "        * reply count\n",
    "        * like count\n",
    "        * quote count\n",
    "        \"\"\"\n",
    "        df = bin_ts[\"eng_dataframe\"]\n",
    "        \n",
    "        rt_count = []\n",
    "        rp_count = []\n",
    "        lc_count = []\n",
    "        qt_count = []\n",
    "        \n",
    "        if len(df)>0:\n",
    "            pms = df[\"tweet_public_metrics\"].tolist()\n",
    "            for pm in pms:\n",
    "                rt_count.append(pm[\"retweet_count\"])\n",
    "                rp_count.append(pm[\"reply_count\"])\n",
    "                lc_count.append(pm[\"like_count\"])\n",
    "                qt_count.append(pm[\"quote_count\"])\n",
    "        \n",
    "        \n",
    "            rt_count = np.sum(rt_count)\n",
    "            rp_count = np.sum(rp_count)\n",
    "            lc_count = np.sum(lc_count)\n",
    "            qt_count = np.sum(qt_count)\n",
    "\n",
    "            pm_vec = [rt_count,rp_count,lc_count,qt_count]\n",
    "            return pm_vec\n",
    "        \n",
    "        else:\n",
    "            pm_vec = [0,0,0,0]\n",
    "            return pm_vec\n",
    "        \n",
    "\n",
    "class NonEngagementPublicMetricsExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        pass\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        public metrics types:\n",
    "        * retweet count\n",
    "        * reply count\n",
    "        * like count\n",
    "        * quote count\n",
    "        \"\"\"\n",
    "        df = bin_ts[\"noeng_dataframe\"]\n",
    "        \n",
    "        rt_count = []\n",
    "        rp_count = []\n",
    "        lc_count = []\n",
    "        qt_count = []\n",
    "        \n",
    "        if len(df)>0:\n",
    "            pms = df[\"tweet_public_metrics\"].tolist()\n",
    "            for pm in pms:\n",
    "                rt_count.append(pm[\"retweet_count\"])\n",
    "                rp_count.append(pm[\"reply_count\"])\n",
    "                lc_count.append(pm[\"like_count\"])\n",
    "                qt_count.append(pm[\"quote_count\"])\n",
    "        \n",
    "        \n",
    "            rt_count = np.sum(rt_count)\n",
    "            rp_count = np.sum(rp_count)\n",
    "            lc_count = np.sum(lc_count)\n",
    "            qt_count = np.sum(qt_count)\n",
    "\n",
    "            pm_vec = [rt_count,rp_count,lc_count,qt_count]\n",
    "            return pm_vec\n",
    "        \n",
    "        else:\n",
    "            pm_vec = [0,0,0,0]\n",
    "            return pm_vec\n",
    "\n",
    "class DRIdentifier(object):\n",
    "    \n",
    "    def __init__(self,news_df):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        ex_news_df = news_df.explode(\"Twitter Handle\").reset_index(drop=True).explode(\"URL\").reset_index(drop=True)\n",
    "        ex_news_twh = ex_news_df[\"Twitter Handle\"].tolist()\n",
    "        self.news_twh = [e for e in ex_news_twh if type(e)==str]\n",
    "        \n",
    "        self.match_pattern = re.compile(r\"(RT\\s{0,}@[a-zA-Z0-9]*)\")\n",
    "        \n",
    "    def check_drt(self,tweet_text):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        matches = self.match_pattern.search(tweet_text)\n",
    "        \n",
    "        if matches != None:\n",
    "            matches = matches.group(0)\n",
    "            acc = matches.replace(\"RT\",\"\").replace(\":\",\"\").replace(\"@\",\"\").strip()\n",
    "\n",
    "            if acc in self.news_twh:\n",
    "                return True\n",
    "            else:\n",
    "                return False\n",
    "        else:\n",
    "            return False\n",
    "    \n",
    "    def identify(self,df):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "    \n",
    "        df[\"direct_retweet\"] = df[\"text\"].apply(lambda x: self.check_drt(x))\n",
    "        \n",
    "        return df\n",
    "\n",
    "class EngagementDRCountsExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self,news_df):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.drt_identifier = DRIdentifier(news_df)\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        df = bin_ts[\"eng_dataframe\"]\n",
    "        drt_count_vec = [0.0]*7\n",
    "        \n",
    "        if len(df)>0:\n",
    "            \n",
    "            df = self.drt_identifier.identify(df)\n",
    "            \n",
    "            masks = [df[\"matched_partisans\"].apply(lambda x: i in x) for i in range(-3,4,1)]\n",
    "            for idx,ps in enumerate(range(-3,4,1)):\n",
    "                \n",
    "                df_ps = df[masks[idx]]\n",
    "                if len(df_ps)>0:\n",
    "                    drt_counts = df_ps[\"direct_retweet\"].value_counts()\n",
    "                    if True in drt_counts:\n",
    "                        drt_count_vec[idx]+= drt_counts[True]\n",
    "        \n",
    "        return drt_count_vec\n",
    "                \n",
    "                \n",
    "\n",
    "class TotalTweetcountExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        pass\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        return bin_ts[\"eng_dataframe\"].shape[0] + bin_ts[\"noeng_dataframe\"].shape[0]\n",
    "\n",
    "class EngagementNSDistExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self,news_df):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.news_df = news_df\n",
    "        self.news_sources = self.news_df[\"Source\"].tolist()\n",
    "        self.ns_map = {k:v for v,k in enumerate(self.news_sources)}\n",
    "        \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        df = bin_ts[\"eng_dataframe\"]\n",
    "        ns_vec = [0.0]*len(self.news_sources)\n",
    "        \n",
    "        if len(df)>0:\n",
    "            \n",
    "            ms = df[\"matched_sources\"].tolist()\n",
    "            ms = flatten(ms)\n",
    "            \n",
    "            for m in ms:\n",
    "                ns_vec[self.ns_map[m]]=1\n",
    "        \n",
    "        return ns_vec\n",
    "\n",
    "\n",
    "def get_hashtags(urow):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    hashtag_regex = r'(?<![#\\w])#(\\w+)'\n",
    "    hast1 = re.findall(hashtag_regex,urow.text)\n",
    "    if len(urow.referenced_text)>0:\n",
    "        hast1+= re.findall(hashtag_regex,urow.referenced_text[0])\n",
    "    \n",
    "    return list(set(hast1))\n",
    "\n",
    "\n",
    "def identity_tok(x):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    x_clean = []\n",
    "    for i in x:\n",
    "        if len(re.sub(r'[^\\x00-\\x7f]',r'', i))>0:\n",
    "            x_clean.append(i)\n",
    "    return x_clean\n",
    "\n",
    "\n",
    "\n",
    "class MentionsGeneralExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self,vectorizer_path):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.vectorizer = pickle.load(open(vectorizer_path,\"rb\"))\n",
    "        self.feature_names = self.vectorizer.get_feature_names()\n",
    "        self.feature_map = {x:i for i,x in enumerate(self.feature_names)}\n",
    "        \n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        * intialize zero vec with vocab size\n",
    "        * get list of mentions from both engagement and non-engagement tweets\n",
    "        * and fill in binary counts\n",
    "        \"\"\"\n",
    "        df_eng = bin_ts[\"eng_dataframe\"]\n",
    "        df_noeng = bin_ts[\"noeng_dataframe\"]\n",
    "        start,end = bin_ts[\"date_range\"]\n",
    "        \n",
    "        men_vec = [0]*len(self.feature_names)\n",
    "        \n",
    "        mentions = df_eng[\"extracted_mentions\"].tolist() + df_noeng[\"extracted_mentions\"].tolist()\n",
    "        \n",
    "        mentions = flatten(mentions)\n",
    "        \n",
    "        if len(mentions)>0:\n",
    "            for m in mentions:\n",
    "                if m in self.feature_map:\n",
    "                    men_vec[self.feature_map[m]]=1\n",
    "\n",
    "        return men_vec\n",
    "\n",
    "\n",
    "\n",
    "class HashtagGeneralExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self,vectorizer_path):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        \n",
    "        self.vectorizer = pickle.load(open(vectorizer_path,\"rb\"))\n",
    "        self.feature_names = self.vectorizer.get_feature_names()\n",
    "        self.feature_map = {x:i for i,x in enumerate(self.feature_names)}\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        * intialize zero vec with vocab size\n",
    "        * get list of hashtags from both engagement and non-engagement tweets\n",
    "        * and fill in counts\n",
    "        \"\"\"\n",
    "        df_eng = bin_ts[\"eng_dataframe\"]\n",
    "        df_noeng = bin_ts[\"noeng_dataframe\"]\n",
    "        start,end = bin_ts[\"date_range\"]\n",
    "        \n",
    "        hash_vec = [0]*len(self.feature_names)\n",
    "        \n",
    "        hashtags = df_eng[\"extracted_hashtags\"].tolist() + df_noeng[\"extracted_hashtags\"].tolist()\n",
    "        \n",
    "        hashtags = flatten(hashtags)\n",
    "        \n",
    "        if len(hashtags)>0:\n",
    "            for h in hashtags:\n",
    "                if h in self.feature_map:\n",
    "                    hash_vec[self.feature_map[h]]=1\n",
    "        \n",
    "        return hash_vec\n",
    "\n",
    "class TopHashtagExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self,tokenizer,topk=100):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.tokenizer = tokenizer\n",
    "        self.topk=topk\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        df_eng = bin_ts[\"eng_dataframe\"]\n",
    "        df_noeng = bin_ts[\"noeng_dataframe\"]\n",
    "        \n",
    "        hashtags = df_eng[\"extracted_hashtags\"].tolist() + df_noeng[\"extracted_hashtags\"].tolist()\n",
    "        \n",
    "        hashtags = flatten(hashtags)\n",
    "        \n",
    "        hashtags_counter = Counter(hashtags)\n",
    "        \n",
    "        top_hash_tups = hashtags_counter.most_common(self.topk)\n",
    "        \n",
    "        top_hashes = []\n",
    "        for tp in top_hash_tups:\n",
    "            top_hashes.append(tp[0])\n",
    "        \n",
    "        if len(top_hashes)==0:\n",
    "            top_hashes=[\"\"]\n",
    "        \n",
    "        tokenized_out = self.tokenizer.batch_encode_plus(top_hashes,\n",
    "                                       padding=\"max_length\",\n",
    "                                       max_length=8,\n",
    "                                       truncation=True,\n",
    "                                       return_tensors='pt',\n",
    "                                       add_special_tokens=False)\n",
    "        \n",
    "        return tokenized_out\n",
    "\n",
    "class TopMentionsExtractor(FeatureExtractor):\n",
    "    \n",
    "    def __init__(self,tokenizer,topk=100):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.tokenizer=tokenizer\n",
    "        self.topk=topk\n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        * intialize zero vec with vocab size\n",
    "        * get list of hashtags from both engagement and non-engagement tweets\n",
    "        * and fill in counts\n",
    "        \"\"\"\n",
    "        df_eng = bin_ts[\"eng_dataframe\"]\n",
    "        df_noeng = bin_ts[\"noeng_dataframe\"]\n",
    "        \n",
    "        mentions = df_eng[\"extracted_mentions\"].tolist() + df_noeng[\"extracted_mentions\"].tolist()\n",
    "        \n",
    "        mentions = flatten(mentions)\n",
    "        \n",
    "        mentions_counter = Counter(mentions)\n",
    "        \n",
    "        top_mens_tups = mentions_counter.most_common(self.topk)\n",
    "        \n",
    "        top_mentions = []\n",
    "        for tp in top_mens_tups:\n",
    "            top_mentions.append(tp[0])\n",
    "        \n",
    "        if len(top_mentions)==0:\n",
    "            top_mentions=[\"\"]\n",
    "        \n",
    "        tokenized_out = self.tokenizer.batch_encode_plus(top_mentions,\n",
    "                                       padding=\"max_length\",\n",
    "                                       max_length=8,\n",
    "                                       truncation=True,\n",
    "                                       return_tensors='pt',\n",
    "                                       add_special_tokens=False)\n",
    "        \n",
    "        return tokenized_out\n",
    "\n",
    "\n",
    "\n",
    "class EngagementTextExtractorT1(FeatureExtractor):\n",
    "    \"\"\"\n",
    "    # * cls token for 25 most recent tweets -> 25,768\n",
    "    # * aggregate on top of this -> avg+min+max\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,tokenizer,num_tweets=25):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.tokenizer = tokenizer\n",
    "        self.num_tweets = num_tweets\n",
    "            \n",
    "    \n",
    "    def get_tweets(self,eng_bin_df):\n",
    "        \"\"\"\n",
    "        case 1: size of eng_bin_df < self.num_tweets\n",
    "        \"\"\"\n",
    "        \n",
    "        if eng_bin_df.shape[0] == 0:\n",
    "            most_recent_tweets = [\"\"]*25\n",
    "            bin_masks = [0]*25\n",
    "            tokenized_out = self.tokenizer.batch_encode_plus(most_recent_tweets,\n",
    "                                           padding=\"max_length\",\n",
    "                                           max_length=20,\n",
    "                                           truncation=True,\n",
    "                                           return_tensors='pt',\n",
    "                                           add_special_tokens=False)\n",
    "            \n",
    "            tokenized_out[\"bin_masks\"] = bin_masks\n",
    "            return tokenized_out\n",
    "            \n",
    "        \n",
    "        else:\n",
    "            eng_bin_df = eng_bin_df.sort_index()\n",
    "            most_recent_tweets = eng_bin_df[\"text\"].tail(self.num_tweets).tolist()\n",
    "            bin_masks = [1]*len(most_recent_tweets)\n",
    "            if len(most_recent_tweets)<self.num_tweets:\n",
    "                most_recent_tweets+= [\"\"]* (self.num_tweets - len(most_recent_tweets))\n",
    "                bin_masks += [0]*(self.num_tweets - len(bin_masks))\n",
    "            \n",
    "            tokenized_out = self.tokenizer.batch_encode_plus(most_recent_tweets,\n",
    "                                           padding=\"max_length\",\n",
    "                                           max_length=20,\n",
    "                                           truncation=True,\n",
    "                                           return_tensors=\"pt\",\n",
    "                                           add_special_tokens=True)\n",
    "            \n",
    "            tokenized_out[\"bin_masks\"] = bin_masks\n",
    "            \n",
    "            return tokenized_out\n",
    "            \n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        df_eng = bin_ts[\"eng_dataframe\"]\n",
    "        df_neng = bin_ts[\"noeng_dataframe\"]\n",
    "        \n",
    "        tokenized_tweets_eng = self.get_tweets(df_eng)\n",
    "        \n",
    "        tokenized_tweets_neng = self.get_tweets(df_neng)\n",
    "        \n",
    "        return tokenized_tweets_eng, tokenized_tweets_neng\n",
    "    \n",
    "\n",
    "\n",
    "class EngagementTextExtractorT2(FeatureExtractor):\n",
    "    \"\"\"\n",
    "    # * combine all tweets into 1 string :\n",
    "    # * Pick cls token\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,tokenizer,num_tweets=25,max_length=512):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.tokenizer = tokenizer\n",
    "        self.num_tweets = num_tweets\n",
    "        self.max_length = max_length\n",
    "            \n",
    "    \n",
    "    def get_tweets(self,eng_bin_df):\n",
    "        \"\"\"\n",
    "        case 1: size of eng_bin_df < self.num_tweets\n",
    "        \"\"\"\n",
    "        \n",
    "        if eng_bin_df.shape[0] == 0:\n",
    "            tokenized_out = self.tokenizer([\"\"],\n",
    "                                           padding=\"max_length\",\n",
    "                                           max_length=512,\n",
    "                                           truncation=True,\n",
    "                                           return_tensors='pt',\n",
    "                                           add_special_tokens=False)\n",
    "            \n",
    "            \n",
    "            return tokenized_out\n",
    "            \n",
    "        \n",
    "        else:\n",
    "            eng_bin_df = eng_bin_df.sort_index()\n",
    "            most_recent_tweets = eng_bin_df[\"text\"].tail(self.num_tweets).tolist()\n",
    "            all_str = \". \".join(most_recent_tweets)\n",
    "            tokenized_out = self.tokenizer(all_str,\n",
    "                                           padding=\"max_length\",\n",
    "                                           max_length=512,\n",
    "                                           truncation=True,\n",
    "                                           return_tensors=\"pt\",\n",
    "                                           add_special_tokens=True)\n",
    "            \n",
    "            return tokenized_out\n",
    "            \n",
    "    \n",
    "    def extract(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        df_eng = bin_ts[\"eng_dataframe\"]\n",
    "        df_neng = bin_ts[\"noeng_dataframe\"]\n",
    "        \n",
    "        tokenized_tweets_eng = self.get_tweets(df_eng)\n",
    "        \n",
    "        tokenized_tweets_neng = self.get_tweets(df_neng)\n",
    "        \n",
    "        return tokenized_tweets_eng, tokenized_tweets_neng"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee32fbb-a972-4c8e-a3b5-500d30872785",
   "metadata": {},
   "outputs": [],
   "source": [
    "def flatten(list_):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    return [a  for x in list_ for a in x] \n",
    "\n",
    "\n",
    "class FEChainExtractor():\n",
    "    \n",
    "    def __init__(self,news_df,ments_vecpath,hash_vecpath,num_eng=10,num_neng=3):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        start = time.time()\n",
    "        # self.tokenizer = AutoTokenizer.from_pretrained(\"/home/kshvaram/roberta_tweet/twitter-roberta-base/\")\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(\"/home/kshvaram/bert_tweet/twhin-bert-base/\")\n",
    "        stop = time.time()\n",
    "        print(f\"Tokenizer loaded in : {round((stop-start)/60,2)} mins\")\n",
    "        \n",
    "        self.__extractor1 = EngagementCountExtractor()\n",
    "        self.__extractor2 = EngagementTypeCountExtractor()\n",
    "        self.__extractor3 = EngagementPublicMetricsExtractor()\n",
    "        self.__extractor4 = NonEngagementPublicMetricsExtractor()\n",
    "        self.__extractor5 = EngagementDRCountsExtractor(news_df)\n",
    "        self.__extractor6 = TotalTweetcountExtractor()\n",
    "        self.__extractor7 = EngagementNSDistExtractor(news_df)\n",
    "        self.__extractor8 = MentionsGeneralExtractor(ments_vecpath)\n",
    "        self.__extractor9 = HashtagGeneralExtractor(hash_vecpath)\n",
    "        self.__extractor10 = EngagementTextExtractorT1(num_tweets=num_eng,tokenizer=self.tokenizer)\n",
    "        self.__extractor11 = EngagementTextExtractorT2(num_tweets=num_eng,max_length=512,tokenizer=self.tokenizer)\n",
    "        self.__extractor12 = TopHashtagExtractor(topk=50,tokenizer=self.tokenizer)\n",
    "        self.__extractor13 = TopMentionsExtractor(topk=50,tokenizer=self.tokenizer)\n",
    "        \n",
    "        self.fe_factory = {}\n",
    "        self.__load_factory__()\n",
    "    \n",
    "    def __load_factory__(self):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.fe_factory = {\"engagement_counts\":self.__extractor1,\n",
    "                           \"engagement_types\":self.__extractor2,\n",
    "                           \"engagement_publicmetrics\":self.__extractor3,\n",
    "                           \"nonengagement_publicmetrics\":self.__extractor4,\n",
    "                           \"engagement_drt\":self.__extractor5,\n",
    "                           \"total_tweet_counts\":self.__extractor6,\n",
    "                           \"engagement_nsdist\": self.__extractor7,\n",
    "                           \"mentions_usertime\":self.__extractor8,\n",
    "                           \"hashtags_usertime\":self.__extractor9,\n",
    "                           \"text_t1\":self.__extractor10,\n",
    "                           \"top_hashes\":self.__extractor12,\n",
    "                           \"top_mentions\":self.__extractor13}\n",
    "    \n",
    "    def extract_feats(self,bin_ts):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        chain = self.fe_factory\n",
    "        \n",
    "        feats = {fe:chain[fe].extract(bin_ts) for fe in chain.keys()}\n",
    "        \n",
    "        return feats\n",
    "\n",
    "\n",
    "class BinUsers(object):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    def __init__(self,user_paths,bin_freq=\"3MS\",start_date='01/01/2014',end_date='9/01/2021'):\n",
    "        \"\"\"\n",
    "        * user_paths : file paths of user timelines\n",
    "        * bin_freq : size of each bin in terms of time\n",
    "        * start_date : which date to start the binning from\n",
    "        * end_date : which date to stop binning\n",
    "        \"\"\"\n",
    "        self.user_paths = user_paths\n",
    "        self.bin_freq = bin_freq\n",
    "        self.start_date = start_date\n",
    "        self.end_date = end_date\n",
    "    \n",
    "    def date_range_init(self):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        index = pd.date_range(start=self.start_date, \n",
    "                              end=self.end_date, \n",
    "                              freq=self.bin_freq,\n",
    "                              inclusive=\"both\",\n",
    "                              normalize=True)\n",
    "\n",
    "        startd=[]\n",
    "        endd=[]\n",
    "\n",
    "        for i in range(1,len(index),1):\n",
    "            startd.append(index[i-1])\n",
    "            endd.append(index[i])\n",
    "\n",
    "        self.df_dates = pd.DataFrame()\n",
    "        self.df_dates[\"start\"]=pd.to_datetime(startd,utc=True)\n",
    "        self.df_dates[\"end\"]=pd.to_datetime(endd,utc=True)\n",
    "        year_difference = int(pd.to_datetime(self.end_date).year) - int(pd.to_datetime(self.start_date).year)\n",
    "        self.df_dates[\"quarter\"] = [1,2,3,4]*year_difference + [1,2,3]\n",
    "        \n",
    "    \n",
    "    def plot_empty_bin_dist(self,avg_bin_dist):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        fig,ax = plt.subplots(1,1,figsize=(3,4))\n",
    "        sns.boxplot(y=avg_bin_dist,\n",
    "                    ax=ax,\n",
    "                    orient=\"v\",\n",
    "                    width=0.3,\n",
    "                    showmeans=True,\n",
    "                    color=\"tab:red\",\n",
    "                    meanprops={\"marker\":\"o\",\n",
    "                               \"markerfacecolor\":\"white\", \n",
    "                               \"markeredgecolor\":\"black\",\n",
    "                              \"markersize\":\"10\"})\n",
    "        ax.set_ylabel(\"Avg Empty Bins Per User\")\n",
    "        ax.set_xlabel(\"Users\")\n",
    "        plt.show()\n",
    "        \n",
    "    \n",
    "    def bin_users(self,feats_extractor):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        all_user_binned_seqs = []\n",
    "        avg_empty_bins=[]\n",
    "        \n",
    "        self.date_range_init()\n",
    "\n",
    "        for up in tqdm.tqdm(self.user_paths):\n",
    "           \n",
    "            df_u = pd.read_pickle(up)\n",
    "            df_u[\"extracted_hashtags\"] = df_u.apply(lambda x: get_hashtags(x),axis=1) \n",
    "            df_u[\"created_at\"] = pd.to_datetime(df_u[\"created_at\"])\n",
    "            df_u = df_u.loc[df_u[\"created_at\"]>=self.start_date]\n",
    "            df_u = df_u.set_index(\"created_at\")\n",
    "\n",
    "            feat_bins = []\n",
    "            empty_c=0\n",
    "            for index,row in self.df_dates.iterrows():\n",
    "                df_bin_dict = {}\n",
    "                df_bin = df_u.loc[(df_u.index>=row[\"start\"]) & (df_u.index< row[\"end\"])]\n",
    "                df_bin_engagement = df_bin.loc[df_bin[\"matched_partisans\"].str.len()>=1]\n",
    "                df_bin_no_engagement = df_bin.loc[df_bin[\"matched_partisans\"].str.len()<=0]\n",
    "                \n",
    "                df_bin_dict[\"eng_dataframe\"] = df_bin_engagement\n",
    "                df_bin_dict[\"noeng_dataframe\"] = df_bin_no_engagement\n",
    "                df_bin_dict[\"quarter\"] = row[\"quarter\"]\n",
    "                df_bin_dict[\"date_range\"] = (row[\"start\"],row[\"end\"])\n",
    "                \n",
    "                feats_dict = feats_extractor.extract_feats(df_bin_dict)\n",
    "                \n",
    "                feats_dict[\"eng_text_t1\"] = feats_dict[\"text_t1\"][0]\n",
    "                # feats_dict[\"eng_text_t2\"] = feats_dict[\"text_t2\"][0]\n",
    "                feats_dict[\"neng_text_t1\"] = feats_dict[\"text_t1\"][1]\n",
    "                # feats_dict[\"neng_text_t2\"] = feats_dict[\"text_t2\"][1]\n",
    "                \n",
    "                # feats_dict[\"eng_text_text_ids\"] = feats_dict[\"text\"][0][\"input_ids\"][0]\n",
    "                # feats_dict[\"eng_text_atten_mask\"] = feats_dict[\"text\"][0][\"attention_mask\"][0]\n",
    "                # feats_dict[\"neng_text_text_ids\"] = feats_dict[\"text\"][1][\"input_ids\"][0]\n",
    "                # feats_dict[\"neng_text_atten_mask\"] = feats_dict[\"text\"][1][\"attention_mask\"][0]\n",
    "                \n",
    "                del feats_dict[\"text_t1\"]\n",
    "                # del feats_dict[\"text_t2\"]\n",
    "        \n",
    "                feats_dict[\"quarter\"] = df_bin_dict[\"quarter\"]\n",
    "                feats_dict[\"date_range\"] = df_bin_dict[\"date_range\"]\n",
    "                feats_dict[\"username\"] = up.split(\"/\")[-1]\n",
    "            \n",
    "                \n",
    "                feat_bins.append(feats_dict)\n",
    "                \n",
    "                \n",
    "                \n",
    "            all_user_binned_seqs.append(feat_bins)\n",
    "            \n",
    "            \n",
    "            \n",
    "        return all_user_binned_seqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4603a41-71f3-49af-aa31-6ffcaa5f8cdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "news_df = pd.read_pickle(\"../../../Data/all_news_sources/all_news_sources.pkl\")\n",
    "\n",
    "feats_extractor = FEChainExtractor(news_df,\n",
    "                                   ments_vecpath=\"vocabs/mentions_general_vectorizer.pk\",\n",
    "                                   hash_vecpath=\"vocabs/hashtags_general_vectorizer.pk\",\n",
    "                                   num_eng=25,num_neng=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f014db9-fca1-4496-9408-5100743b84cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "users_binned  =  BinUsers(filtered_user_paths,\n",
    "                          bin_freq=\"3MS\",\n",
    "                          start_date='01/01/2015',\n",
    "                          end_date='10/01/2021').bin_users(feats_extractor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a59b6dc-1111-4878-96b2-ad61ce803511",
   "metadata": {},
   "source": [
    "## Train, Val, Test Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c361f3fb-c1af-472d-8bd0-2d0822a36c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def trunc_datetime(someDate):\n",
    "    # remove day, time, timezone info for comparison\n",
    "    return someDate.replace(day=1, hour=0, minute=0, second=0, microsecond=0).tz_localize(None)\n",
    "\n",
    "def split_all_user_sequences(all_sequences):\n",
    "    \"\"\"\n",
    "    * all_sequences: list of list of dicts , where each list item represents a user\n",
    "    * and the user list contains dicts wher each of them represents a bin/timestep\n",
    "    * returns splits of user sequences binned according to 3 year time period\n",
    "    \"\"\"\n",
    "    dates = {1:(trunc_datetime(pd.to_datetime(\"2015-01-01\",format='%Y-%M-%d')),\n",
    "                trunc_datetime(pd.to_datetime(\"2019-01-01\",format='%Y-%M-%d'))),\n",
    "             2:(trunc_datetime(pd.to_datetime(\"2016-01-01\",format='%Y-%M-%d')),\n",
    "                trunc_datetime(pd.to_datetime(\"2020-01-01\",format='%Y-%M-%d'))),\n",
    "             3:(trunc_datetime(pd.to_datetime(\"2017-01-01\",format='%Y-%M-%d')),\n",
    "                trunc_datetime(pd.to_datetime(\"2021-01-01\",format='%Y-%M-%d'))),\n",
    "             4:(trunc_datetime(pd.to_datetime(\"2018-01-01\",format='%Y-%M-%d')),\n",
    "                trunc_datetime(pd.to_datetime(\"2022-01-01\",format='%Y-%M-%d')))}\n",
    "    \n",
    "    run1_obs = []\n",
    "    run2_obs = []\n",
    "    run3_obs = []\n",
    "    run4_obs = []\n",
    "    \n",
    "    \n",
    "    for user_sequence in tqdm.tqdm(all_sequences):\n",
    "        \n",
    "        binned_sequence_r1 = []\n",
    "        binned_sequence_r2 = []\n",
    "        binned_sequence_r3 = []\n",
    "        binned_sequence_r4 = []\n",
    "        \n",
    "        for user_bin_dict in user_sequence:\n",
    "            \n",
    "            start_date = trunc_datetime(user_bin_dict[\"date_range\"][0])\n",
    "            end_date = trunc_datetime(user_bin_dict[\"date_range\"][1])\n",
    "            \n",
    "            if start_date >= dates[1][0] and end_date <= dates[1][1]:\n",
    "                binned_sequence_r1.append(user_bin_dict)\n",
    "                \n",
    "            if start_date >= dates[2][0] and end_date <= dates[2][1]:\n",
    "                binned_sequence_r2.append(user_bin_dict)\n",
    "            \n",
    "            if start_date >= dates[3][0] and end_date <= dates[3][1]:\n",
    "                binned_sequence_r3.append(user_bin_dict)\n",
    "            \n",
    "            if start_date >= dates[4][0] and end_date <= dates[4][1]:\n",
    "                binned_sequence_r4.append(user_bin_dict)\n",
    "            \n",
    "        \n",
    "        assert len(binned_sequence_r1) == len(binned_sequence_r2) == len(binned_sequence_r3)\n",
    "        assert len(binned_sequence_r1)-1 ==len(binned_sequence_r4)\n",
    "        \n",
    "        run1_obs.append(binned_sequence_r1)\n",
    "        run2_obs.append(binned_sequence_r2)\n",
    "        run3_obs.append(binned_sequence_r3)\n",
    "        run4_obs.append(binned_sequence_r4)\n",
    "        \n",
    "    \n",
    "    return run1_obs, run2_obs, run3_obs, run4_obs\n",
    "\n",
    "def use_sliding_window(seq,window_size=9):\n",
    "    \"\"\"\n",
    "    Seq : Indices of the timestep bins\n",
    "    \"\"\"\n",
    "    sub_sequences = []\n",
    "    \n",
    "    for w in range(len(seq)-window_size+1):\n",
    "        sub_sequences.append(seq[w:w+window_size])\n",
    "        \n",
    "    return sub_sequences\n",
    "\n",
    "\n",
    "def get_train_val_test_split(user_sequences,train_val_split=[0.8,0.2]):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    test = [useq[4:] for useq in user_sequences]\n",
    "    \n",
    "    num_users = len(user_sequences)\n",
    "    \n",
    "    train_size = int(train_val_split[0]*num_users)\n",
    "    val_size = int(train_val_split[1]*num_users)\n",
    "    \n",
    "    random.seed(42)\n",
    "    train_user_indices = random.sample([i for i in range(num_users)],train_size)\n",
    "    val_user_indices = [j for j in range(num_users) if j not in  train_user_indices] \n",
    "    \n",
    "    train = [user_sequences[i][:4] for i in train_user_indices]\n",
    "    \n",
    "    val = [user_sequences[i][:4] for i in val_user_indices]\n",
    "    \n",
    "    print(f\"No of users in Train : {len(train)}\")\n",
    "    print(f\"No of users in Val : {len(val)}\")\n",
    "    print(f\"No of users in Test : {len(test)}\\n\")\n",
    "    \n",
    "    return train, val, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c79c9cab-f030-42c3-b733-99ee42c50c6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "run1_split, run2_split, run3_split, run4_split = split_all_user_sequences(users_binned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69ed2477-ca6a-4c56-9a84-46981353dff2",
   "metadata": {},
   "outputs": [],
   "source": [
    "run1_split_sw = [use_sliding_window(useq,window_size=9) for useq in run1_split]\n",
    "run2_split_sw = [use_sliding_window(useq,window_size=9) for useq in run2_split]\n",
    "run3_split_sw = [use_sliding_window(useq,window_size=9) for useq in run3_split]\n",
    "run4_split_sw = [use_sliding_window(useq,window_size=9) for useq in run4_split]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7171ce4-c524-48c3-950b-c310873ba46c",
   "metadata": {},
   "outputs": [],
   "source": [
    "r1_train, r1_val, r1_test = get_train_val_test_split(run1_split_sw,train_val_split=[0.8,0.2])\n",
    "r2_train, r2_val, r2_test = get_train_val_test_split(run2_split_sw,train_val_split=[0.8,0.2])\n",
    "r3_train, r3_val, r3_test = get_train_val_test_split(run3_split_sw,train_val_split=[0.8,0.2])\n",
    "r4_train, r4_val, r4_test = get_train_val_test_split(run4_split_sw,train_val_split=[0.8,0.2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6525193a-52e1-4ca2-b3c4-8af9a7ed0d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "r1_train = flatten(r1_train) \n",
    "r1_val = flatten(r1_val)\n",
    "r1_test = flatten(r1_test)\n",
    "\n",
    "r2_train = flatten(r2_train) \n",
    "r2_val = flatten(r2_val)\n",
    "r2_test = flatten(r2_test)\n",
    "\n",
    "r3_train = flatten(r3_train) \n",
    "r3_val = flatten(r3_val)\n",
    "r3_test = flatten(r3_test)\n",
    "\n",
    "r4_train = flatten(r4_train) \n",
    "r4_val = flatten(r4_val)\n",
    "r4_test = flatten(r4_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a68d8cb-5198-4d07-8f8b-54041ac0e9b2",
   "metadata": {},
   "source": [
    "## Postprocessing Train, Val, Test - Precomputing Features and saving to Disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b04784b8-e3c7-4880-bee7-daa9bcab35a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def zscore_obs(sequence_counts):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    binned_mat = np.array(sequence_counts)[:-1]\n",
    "    means = np.mean(binned_mat,axis=0)\n",
    "    std_devs = np.std(binned_mat,axis=0)\n",
    "\n",
    "    normalized_seq = (np.array(sequence_counts) - means)/std_devs\n",
    "    normalized_seq[np.isnan(normalized_seq)]=0\n",
    "    normalized_seq[np.isinf(normalized_seq)]=0\n",
    "\n",
    "    return normalized_seq , means, std_devs\n",
    "\n",
    "def one_hot_encode_quarters(quarter_ids):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    quarters = []\n",
    "    \n",
    "    for q in quarter_ids:\n",
    "        quarter_encoded = [0,0,0,0]\n",
    "        quarter_encoded[q-1]=1\n",
    "        quarters.append(quarter_encoded)\n",
    "        \n",
    "    return quarters\n",
    "\n",
    "def combine_across_time(sequence_mat):\n",
    "    \"\"\"\n",
    "    (timesteps,feature dim) -> (1, feature dim)\n",
    "    \n",
    "    * Perform sum and convert values >0 to 1 and <=0 to 0\n",
    "    * fill nans and infs\n",
    "    \"\"\"\n",
    "    agg_mat = np.sum(sequence_mat[:-1,:],axis=0)\n",
    "    return (agg_mat > 1).astype(int)\n",
    "\n",
    "def get_bin_masks(text_sequence):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    bin_masks = []\n",
    "    for i in text_sequence.shape[0]:\n",
    "        \n",
    "        if torch.sum(text_sequence[i,:])==text_sequence.shape[-1]:\n",
    "            bin_masks.append(True)\n",
    "        else:\n",
    "            bin_masks.append(False)\n",
    "    \n",
    "    return torch.Tensor(bin_masks)\n",
    "\n",
    "def remove_empty_bins(text_sequence,attnmask):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    filtered_text = []\n",
    "    filtered_attn = []\n",
    "    for i in text_sequence.shape[0]:\n",
    "        \n",
    "        if not torch.sum(text_sequence[i,:])==text_sequence.shape[-1]:\n",
    "            \n",
    "            filtered_text.append(text_sequence[i,:].unsqueeze())\n",
    "            filtered_attn.append(attnmask[i,:].unsequeeze())\n",
    "    \n",
    "    if len(filtered_text)>0:\n",
    "        filtered_text = torch.concat(filtered_text,dim=0)\n",
    "        filtered_attn = torch.concat(filtered_attn,dim=0)\n",
    "    \n",
    "    if len(filtered_text)<0:\n",
    "        pass\n",
    "    \n",
    "    return filtered_text, filtered_attn\n",
    "    \n",
    "    \n",
    "def get_reps(eng_text, eng_attn, model,dim=1):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    out = None\n",
    "    \n",
    "    \n",
    "    with torch.set_grad_enabled(False):\n",
    "        \n",
    "        out = model(eng_text,eng_attn)\n",
    "        \n",
    "    return out.last_hidden_state\n",
    "    \n",
    "    \n",
    "def get_text_reps_cls_text(tokens,attn,bin_masks,model):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    pass\n",
    "    \n",
    "def get_text_reps_cls(tokens,attn,model,ttype=1):\n",
    "    \"\"\"\n",
    "    type 1: \n",
    "    * Select cls tokens from each tweet rep\n",
    "    * agg these cls tokens\n",
    "    * eg: [8,x,20]->[8,x,20,768] -> [8,x,1,768] -> ideally [8,1,768]\n",
    "    \n",
    "    * the above can also be done as follows:\n",
    "    * for every tweet x -> concatenate cls token + avg token embeddings\n",
    "    * now aggregate this resulting tensor -> sum,min,max,mean etc ...\n",
    "    \n",
    "    type 2:\n",
    "    * select cls token rep \n",
    "    * [8,1,512] -> [8,512,768] -> [8,768]\n",
    "    * \n",
    "    \"\"\"\n",
    "    res_dict = {}\n",
    "    hidden_rep = None\n",
    "    \n",
    "    \n",
    "    if ttype==1:\n",
    "        \n",
    "        tokens = tokens[:-1,:,:]\n",
    "        attn = attn[:-1,:,:]\n",
    "        \n",
    "        num_timesteps = tokens.shape[0]\n",
    "        num_tweets = tokens.shape[1]\n",
    "        max_tweetlength = tokens.shape[2]\n",
    "        tokens = torch.reshape(tokens,(num_timesteps*num_tweets,max_tweetlength))\n",
    "        attn = torch.reshape(attn,(num_timesteps*num_tweets,max_tweetlength))\n",
    "        \n",
    "        acc = [] # for loop since 9 * 25 is too large a batch for the current gpu\n",
    "        \n",
    "#         for i in range(9):\n",
    "        with torch.set_grad_enabled(False):\n",
    "            out = model(tokens,attn)\n",
    "            hidden_rep = out.last_hidden_state\n",
    "        \n",
    "        \n",
    "            \n",
    "        hidden_rep = torch.reshape(hidden_rep,(num_timesteps,num_tweets,max_tweetlength,768))\n",
    "        cls_reps = hidden_rep[:,:,0,:] # select cls token -> (timesteps,num_tweets,768)\n",
    "        # padding_ind = bin_masks.index(0)\n",
    "        other_token_reps = hidden_rep[:,:,1:,:] # select other tokens -> (timesteps,num_tweets,24,768)\n",
    "        avg_token_reps = torch.mean(other_token_reps,dim=2) # avg other tokens -> (timesteps,num_tweets,768)\n",
    "        fused_reps = torch.concat([cls_reps,avg_token_reps],dim=-1) # concat -> (timesteps,num_tweets,1536)\n",
    "        \n",
    "        \n",
    "        masked_hidden_rep = fused_reps\n",
    "        hr_mean = torch.mean(masked_hidden_rep,dim=1)\n",
    "        res_dict[\"mean\"] = hr_mean.detach().cpu()\n",
    "        \n",
    "    \n",
    "    else:\n",
    "        tokens = tokens[:-1,:]\n",
    "        attn = attn[:-1,:]\n",
    "        \n",
    "        with torch.set_grad_enabled(False):\n",
    "            out = model(tokens,attn)\n",
    "            hidden_rep = out.last_hidden_state\n",
    "        \n",
    "        # type 2 where all tweets are concatenated into one string of length 512 (timesteps,512,768)\n",
    "        cls_rep = hidden_rep[:,0,:] # (9,768)\n",
    "        other_token_reps = hidden_rep[:,1:,:] #(9,511,768)\n",
    "        \n",
    "        res_dict[\"cls\"] = cls_rep.detach().cpu()\n",
    "        res_dict[\"sum\"] = torch.sum(other_token_reps,dim=1).detach().cpu()\n",
    "        res_dict[\"min\"] = torch.min(other_token_reps,dim=1)[0].detach().cpu()\n",
    "        res_dict[\"max\"] = torch.max(other_token_reps,dim=1)[0].detach().cpu()\n",
    "        res_dict[\"mean\"] = torch.mean(other_token_reps,dim=1).detach().cpu()\n",
    "        \n",
    "    \n",
    "    return res_dict\n",
    "\n",
    "def get_hash_mens_reps(tokens,attn,model):\n",
    "    \"\"\"\n",
    "    tokens shape : 50,8\n",
    "    attn shape : 50,8\n",
    "    \n",
    "    # 1st level aggregation\n",
    "    \n",
    "    out[:,0,:] -> cls (50,768)\n",
    "    torch.mean(out[:,1:,:]) -> (50,768)\n",
    "    \n",
    "    torch.concat() -> 50,1536\n",
    "    torch.mean -> 1, 1536\n",
    "    \n",
    "    \n",
    "    \"\"\"\n",
    "    # aggregation -> \n",
    "    with torch.set_grad_enabled(False):\n",
    "        \n",
    "        \n",
    "        out = model(tokens,attn)\n",
    "        hidden_rep = out.last_hidden_state\n",
    "\n",
    "        cls_rep = hidden_rep[:,0,:] \n",
    "        other_token_reps = hidden_rep[:,1:,:]\n",
    "\n",
    "        otr_mean = torch.mean(other_token_reps,dim=1)\n",
    "\n",
    "        cls_otr = torch.concat([cls_rep,otr_mean],dim=-1)\n",
    "\n",
    "        cls_otr_mean = torch.mean(cls_otr,dim=0)\n",
    "\n",
    "        return cls_otr_mean\n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "def extract_sequence_feats(dataset):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    \n",
    "    seq_feats = []\n",
    "    \n",
    "    # model = AutoModel.from_pretrained(\"/home/kshvaram/roberta_tweet/twitter-roberta-base/\")\n",
    "    model = AutoModel.from_pretrained(\"/home/kshvaram/bert_tweet/twhin-bert-base/\")\n",
    "    \n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "    model= nn.DataParallel(model)\n",
    "    model.to(device)\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    for seqd in tqdm.tqdm(dataset):\n",
    "        \n",
    "        eng_counts = []\n",
    "        eng_type_counts = []\n",
    "        eng_pm_counts = []\n",
    "        neng_pm_counts = []\n",
    "        eng_drt_counts = []\n",
    "        tt_counts = []\n",
    "        eng_ns = []\n",
    "        mentions = []\n",
    "        hashtags = []\n",
    "\n",
    "        quarter_ids = []\n",
    "\n",
    "        eng1_text = []\n",
    "        eng1_attn = []\n",
    "        eng1_bin_masks = []\n",
    "        \n",
    "        eng2_text = []\n",
    "        eng2_attn = []\n",
    "        # eng2_bin_masks = []\n",
    "        \n",
    "        \n",
    "        neng1_text = []\n",
    "        neng1_attn = []\n",
    "        neng1_bin_masks = []\n",
    "        \n",
    "        neng2_text = []\n",
    "        neng2_attn = []\n",
    "        # neng2_bin_masks = []\n",
    "        \n",
    "        date_range = []\n",
    "        usernames = []\n",
    "        \n",
    "        top_mens_tokens = []\n",
    "        top_hash_tokens = []\n",
    "        top_mens_attn = []\n",
    "        top_hash_attn = []\n",
    "        \n",
    "        \n",
    "        for obs in seqd:\n",
    "        \n",
    "            eng_counts.append(obs[\"engagement_counts\"])\n",
    "            eng_type_counts.append(obs[\"engagement_types\"])\n",
    "            eng_pm_counts.append(obs[\"engagement_publicmetrics\"])\n",
    "            neng_pm_counts.append(obs[\"nonengagement_publicmetrics\"])\n",
    "            eng_drt_counts.append(obs[\"engagement_drt\"])\n",
    "            tt_counts.append(obs[\"total_tweet_counts\"])\n",
    "            eng_ns.append(obs[\"engagement_nsdist\"])\n",
    "            mentions.append(obs[\"mentions_usertime\"])\n",
    "            hashtags.append(obs[\"hashtags_usertime\"])\n",
    "            quarter_ids.append(obs[\"quarter\"])\n",
    "            date_range.append(obs[\"date_range\"])\n",
    "            usernames.append(obs[\"username\"])\n",
    "            \n",
    "            eng1_text.append(obs[\"eng_text_t1\"][\"input_ids\"].unsqueeze(0))\n",
    "            eng1_attn.append(obs[\"eng_text_t1\"][\"attention_mask\"].unsqueeze(0))\n",
    "            eng1_bin_masks.append(obs[\"eng_text_t1\"][\"bin_masks\"])\n",
    "            \n",
    "            \n",
    "            neng1_text.append(obs[\"neng_text_t1\"][\"input_ids\"].unsqueeze(0))\n",
    "            neng1_attn.append(obs[\"neng_text_t1\"][\"attention_mask\"].unsqueeze(0))\n",
    "            neng1_bin_masks.append(obs[\"neng_text_t1\"][\"bin_masks\"])\n",
    "            \n",
    "            top_mens_tokens.append(obs[\"top_mentions\"][\"input_ids\"].unsqueeze(0))\n",
    "            top_hash_tokens.append(obs[\"top_hashes\"][\"input_ids\"].unsqueeze(0))\n",
    "            top_mens_attn.append(obs[\"top_mentions\"][\"attention_mask\"].unsqueeze(0))\n",
    "            top_hash_attn.append(obs[\"top_hashes\"][\"attention_mask\"].unsqueeze(0))\n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "        \n",
    "        # engagement counts - 1x7\n",
    "        eng_counts_norm, c_means, c_std_dev = zscore_obs(eng_counts)\n",
    "        \n",
    "        # engagement type counts - url,mentions,both - 1x3\n",
    "        eng_type_norm,_,_ = zscore_obs(eng_type_counts)\n",
    "        \n",
    "        # engagement and non-engagement public metrics - retweeted,liked,shared,quote - 1x4\n",
    "        eng_pm_norm,_,_ = zscore_obs(eng_pm_counts)\n",
    "        neng_pm_norm,_,_ = zscore_obs(neng_pm_counts)\n",
    "        \n",
    "        # total_tweet_count in timestep bin - 1x1\n",
    "        tt_counts_norm,_,_ = zscore_obs(tt_counts)\n",
    "        tt_counts_norm = tt_counts_norm.reshape(-1,1)\n",
    "        \n",
    "        # engagement drt counts per stance - 1x7\n",
    "        eng_drt_counts_norm,_,_ = zscore_obs(eng_drt_counts)\n",
    "        \n",
    "        # engaged news sources - binary vec - 1x522\n",
    "        eng_ns = np.array(eng_ns)\n",
    "        \n",
    "        # top mentioned user accounts - binary vec - 1x5000\n",
    "        mentions = np.array(mentions)\n",
    "        \n",
    "        # top used hashtags - binary vec - 1x5000\n",
    "        hashtags = np.array(hashtags)\n",
    "        \n",
    "        # time features - quarter one hot encoded - 1x4\n",
    "        quarter_ids_enc = one_hot_encode_quarters(quarter_ids)\n",
    "        quarter_ids_enc = np.array(quarter_ids_enc)\n",
    "        quarter_ids = np.array(quarter_ids)\n",
    "        \n",
    "        \n",
    "        # text representations\n",
    "        eng1_text = torch.concat(eng1_text,dim=0).to(device)\n",
    "        eng1_attn = torch.concat(eng1_attn,dim=0).to(device)\n",
    "        \n",
    "        neng1_text = torch.concat(neng1_text,dim=0).to(device)\n",
    "        neng1_attn = torch.concat(neng1_attn,dim=0).to(device)\n",
    "        \n",
    "        eng1 = get_text_reps_cls(eng1_text,eng1_attn,model,ttype=1)\n",
    "        neng1 = get_text_reps_cls(neng1_text,neng1_attn,model,ttype=1)\n",
    "        \n",
    "        \n",
    "        # top hashes, mentions\n",
    "        \n",
    "        top_hashes_toks = top_hash_tokens[-2].to(device)\n",
    "        top_hashes_attn = top_hash_attn[-2].to(device)\n",
    "        top_mens_toks = top_mens_tokens[-2].to(device)\n",
    "        top_mens_attn = top_mens_attn[-2].to(device)\n",
    "        \n",
    "        \n",
    "        # last time step of the input sequence\n",
    "        top_hash_reps = get_hash_mens_reps(top_hashes_toks.squeeze(0),\n",
    "                                           top_hashes_attn.squeeze(0),\n",
    "                                           model)\n",
    "        \n",
    "        top_mens_reps = get_hash_mens_reps(top_mens_toks.squeeze(0),\n",
    "                                           top_mens_attn.squeeze(0),\n",
    "                                           model)\n",
    "        \n",
    "        \n",
    "        \n",
    "        eng1_text = None\n",
    "        eng1_attn = None\n",
    "        eng2_text = None\n",
    "        eng2_attn = None\n",
    "        neng1_text = None\n",
    "        neng1_attn = None\n",
    "        neng2_text = None\n",
    "        neng2_attn = None\n",
    "        top_hashes_toks = None\n",
    "        top_hashes_attn = None\n",
    "        top_mens_toks = None\n",
    "        top_mens_attn = None\n",
    "        \n",
    "        del eng1_text\n",
    "        del eng1_attn\n",
    "        del eng2_text\n",
    "        del eng2_attn\n",
    "        del neng1_text\n",
    "        del neng1_attn\n",
    "        del neng2_text\n",
    "        del neng2_attn\n",
    "        del top_hashes_toks\n",
    "        del top_hashes_attn\n",
    "        del top_mens_toks\n",
    "        del top_mens_attn\n",
    "        \n",
    "        \n",
    "        \n",
    "        output_dict = {\"eng_counts\":eng_counts_norm[:-1,:].astype(np.float16),\n",
    "                       \"eng_type\":eng_type_norm[:-1,:].astype(np.float16),\n",
    "                       \"eng_pm\":eng_pm_norm[:-1,:].astype(np.float16),\n",
    "                       \"neng_pm\":neng_pm_norm[:-1,:].astype(np.float16),\n",
    "                       \"total_count\":tt_counts_norm[:-1,:].astype(np.float16),\n",
    "                       \"drt_count\":eng_drt_counts_norm[:-1,:].astype(np.float16),\n",
    "                       \"ns_feat\":eng_ns[:-1,:].astype(np.float16),\n",
    "                       \"mentions_feat\":mentions[:-1,:].astype(np.float16),\n",
    "                       \"hashtag_feat\":hashtags[:-1,:].astype(np.float16),\n",
    "                       \"engagement_text_feats_1\":eng1,\n",
    "                       \"nengagement_text_feats_1\":neng1,\n",
    "                       \"tophashes\":top_hash_reps,\n",
    "                       \"topmentions\":top_mens_reps,\n",
    "                       \"input_time_feat\":quarter_ids_enc[:-1,:].astype(np.float16),\n",
    "                       \"output_time_feat\":quarter_ids_enc[-1,:].astype(np.float16),\n",
    "                       \"output_quarter\":quarter_ids[-1],\n",
    "                       \"label_zscored\":eng_counts_norm[-1,:].astype(np.float16),\n",
    "                       \"label_original\":np.array(eng_counts)[-1,:].astype(np.float16),\n",
    "                       \"means\":c_means.astype(np.float16),\n",
    "                       \"stddevs\":c_std_dev.astype(np.float16),\n",
    "                       \"original_counts\":np.array(eng_counts),\n",
    "                       \"date_ranges\":date_range,\n",
    "                       \"username\":usernames}\n",
    "        \n",
    "        seq_feats.append(output_dict)\n",
    "        \n",
    "    return seq_feats\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b39e2af1-2efa-47ea-88b0-d58ba3f56790",
   "metadata": {},
   "outputs": [],
   "source": [
    "r1_train = extract_sequence_feats(r1_train) \n",
    "\n",
    "with open(\"pickled_data/r1_train.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r1_train,wp)\n",
    "\n",
    "r1_train = None\n",
    "\n",
    "del r1_train\n",
    "    \n",
    "r1_val = extract_sequence_feats(r1_val)\n",
    "\n",
    "with open(\"pickled_data/r1_val.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r1_val,wp)\n",
    "    \n",
    "r1_val = None\n",
    "\n",
    "del r1_val\n",
    "\n",
    "r1_test = extract_sequence_feats(r1_test)\n",
    "\n",
    "with open(\"pickled_data/r1_test.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r1_test,wp)\n",
    "\n",
    "r1_test = None\n",
    "\n",
    "del r1_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e20a810-cd86-477e-bcab-ad5ec829e1c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "r2_train = extract_sequence_feats(r2_train) \n",
    "\n",
    "with open(\"pickled_data/r2_train.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r2_train,wp)\n",
    "    \n",
    "r2_val = extract_sequence_feats(r2_val)\n",
    "\n",
    "\n",
    "with open(\"pickled_data/r2_val.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r2_val,wp)\n",
    "\n",
    "r2_test = extract_sequence_feats(r2_test)\n",
    "\n",
    "\n",
    "with open(\"pickled_data/r2_test.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r2_test,wp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4f85a14-9bc1-4f2e-a1fe-ee8831453712",
   "metadata": {},
   "outputs": [],
   "source": [
    "r3_train = extract_sequence_feats(r3_train) \n",
    "\n",
    "with open(\"pickled_data/r3_train.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r3_train,wp)\n",
    "    \n",
    "r3_val = extract_sequence_feats(r3_val)\n",
    "\n",
    "with open(\"pickled_data/r3_val.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r3_val,wp)\n",
    "    \n",
    "r3_test = extract_sequence_feats(r3_test)\n",
    "\n",
    "with open(\"pickled_data/r3_test.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r3_test,wp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c7b5591-f7a8-4686-9fe3-af2700fb1918",
   "metadata": {},
   "outputs": [],
   "source": [
    "r4_train = extract_sequence_feats(r4_train) \n",
    "\n",
    "with open(\"pickled_data/r4_train.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r4_train,wp)\n",
    "    \n",
    "r4_val = extract_sequence_feats(r4_val)\n",
    "\n",
    "with open(\"pickled_data/r4_val.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r4_val,wp)\n",
    "    \n",
    "r4_test = extract_sequence_feats(r4_test)\n",
    "    \n",
    "with open(\"pickled_data/r4_test.pkl\",\"wb\") as wp:\n",
    "    pickle.dump(r4_test,wp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}