karaolidis/qrust
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add finbert sentiment analysis

Browse Source 
Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>
This commit is contained in:
Nick Karaolidis
2024-02-03 18:58:40 +00:00
parent 973917dad2
commit 65c9ae8b25
26 changed files with 31460 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/config.rs
View File

@@ -4,7 +4,19 @@ use reqwest::{
header::{HeaderMap, HeaderName, HeaderValue},
Client,
};
use std::{env, num::NonZeroU32, sync::Arc};
use rust_bert::{
pipelines::{
common::{ModelResource, ModelType},
sequence_classification::{SequenceClassificationConfig, SequenceClassificationModel},
},
resources::LocalResource,
};
use std::{
env,
num::NonZeroU32,
path::PathBuf,
sync::{Arc, Mutex},
};
pub const ALPACA_ASSET_API_URL: &str = "https://api.alpaca.markets/v2/assets";
pub const ALPACA_CLOCK_API_URL: &str = "https://api.alpaca.markets/v2/clock";
@@ -23,6 +35,8 @@ pub struct Config {
pub alpaca_rate_limit: DefaultDirectRateLimiter,
pub alpaca_source: Source,
pub clickhouse_client: clickhouse::Client,
pub max_bert_inputs: usize,
pub sequence_classifier: Arc<Mutex<SequenceClassificationModel>>,
}
impl Config {
@@ -41,6 +55,11 @@ impl Config {
env::var("CLICKHOUSE_PASSWORD").expect("CLICKHOUSE_PASSWORD must be set.");
let clickhouse_db = env::var("CLICKHOUSE_DB").expect("CLICKHOUSE_DB must be set.");
let max_bert_inputs: usize = env::var("MAX_BERT_INPUTS")
.expect("MAX_BERT_INPUTS must be set.")
.parse()
.expect("MAX_BERT_INPUTS must be a positive integer.");
Self {
alpaca_client: Client::builder()
.default_headers(HeaderMap::from_iter([
@@ -69,6 +88,26 @@ impl Config {
.with_database(clickhouse_db),
alpaca_api_key,
alpaca_api_secret,
max_bert_inputs,
sequence_classifier: Arc::new(Mutex::new(
SequenceClassificationModel::new(SequenceClassificationConfig::new(
ModelType::Bert,
ModelResource::Torch(Box::new(LocalResource {
local_path: PathBuf::from("./models/finbert/rust_model.ot"),
})),
LocalResource {
local_path: PathBuf::from("./models/finbert/config.json"),
},
LocalResource {
local_path: PathBuf::from("./models/finbert/vocab.txt"),
},
None,
true,
None,
None,
))
.unwrap(),
)),
}
}

16
src/database/news.rs
View File

@@ -25,25 +25,31 @@ pub async fn delete_where_symbols<T>(clickhouse_client: &Client, symbols: &[T])
where
T: AsRef<str> + Serialize + Send + Sync,
{
let remaining_symbols = assets::select(clickhouse_client)
.await
.into_iter()
.map(|asset| asset.abbreviation)
.collect::<Vec<_>>();
clickhouse_client
.query("DELETE FROM news WHERE hasAny(symbols, ?)")
.query("DELETE FROM news WHERE hasAny(symbols, ?) AND NOT hasAny(symbols, ?)")
.bind(symbols)
.bind(remaining_symbols)
.execute()
.await
.unwrap();
}
pub async fn cleanup(clickhouse_client: &Client) {
let assets = assets::select(clickhouse_client).await;
let symbols = assets
let remaining_symbols = assets::select(clickhouse_client)
.await
.into_iter()
.map(|asset| asset.abbreviation)
.collect::<Vec<_>>();
clickhouse_client
.query("DELETE FROM news WHERE NOT hasAny(symbols, ?)")
.bind(symbols)
.bind(remaining_symbols)
.execute()
.await
.unwrap();

50
src/threads/data/backfill.rs
View File

@@ -4,18 +4,19 @@ use crate::{
database,
types::{
alpaca::{api, Source},
news::Prediction,
Asset, Bar, Class, News, Subset,
},
utils::{duration_until, last_minute, FIFTEEN_MINUTES, ONE_MINUTE},
};
use backoff::{future::retry, ExponentialBackoff};
use log::{error, info};
use log::{error, info, warn};
use std::{collections::HashMap, sync::Arc};
use time::OffsetDateTime;
use tokio::{
join, spawn,
sync::{mpsc, oneshot, Mutex, RwLock},
task::JoinHandle,
task::{spawn_blocking, JoinHandle},
time::sleep,
};
@@ -103,11 +104,14 @@ async fn handle_backfill_message(
match message.action {
Action::Backfill => {
for symbol in symbols {
if let Some(job) = backfill_jobs.remove(&symbol) {
if let Some(job) = backfill_jobs.get(&symbol) {
if !job.is_finished() {
job.abort();
warn!(
"{:?} - Backfill for {} is already running, skipping.",
thread_type, symbol
);
continue;
}
let _ = job.await;
}
let app_config = app_config.clone();
@@ -361,7 +365,41 @@ async fn execute_backfill_news(
return;
}
let backfill = (news.last().unwrap().clone(), symbol.clone()).into();
let app_config_clone = app_config.clone();
let inputs = news
.iter()
.map(|news| format!("{}\n\n{}", news.headline, news.content))
.collect::<Vec<_>>();
let predictions: Vec<Prediction> = spawn_blocking(move || {
inputs
.chunks(app_config_clone.max_bert_inputs)
.flat_map(|inputs| {
app_config_clone
.sequence_classifier
.lock()
.unwrap()
.predict(inputs.iter().map(String::as_str).collect::<Vec<_>>())
.into_iter()
.map(|label| Prediction::try_from(label).unwrap())
.collect::<Vec<_>>()
})
.collect()
})
.await
.unwrap();
let news = news
.into_iter()
.zip(predictions.into_iter())
.map(|(news, prediction)| News {
sentiment: prediction.sentiment,
confidence: prediction.confidence,
..news
})
.collect::<Vec<_>>();
let backfill = (news[0].clone(), symbol.clone()).into();
database::news::upsert_batch(&app_config.clickhouse_client, news).await;
database::backfills::upsert(&app_config.clickhouse_client, &thread_type, &backfill).await;

16
src/threads/data/mod.rs
View File

@@ -27,16 +27,6 @@ pub struct Guard {
pub pending_unsubscriptions: HashMap<String, Asset>,
}
impl Guard {
pub fn new() -> Self {
Self {
symbols: HashSet::new(),
pending_subscriptions: HashMap::new(),
pending_unsubscriptions: HashMap::new(),
}
}
}
#[derive(Clone, Copy, Debug)]
pub enum ThreadType {
Bars(Class),
@@ -86,7 +76,11 @@ async fn init_thread(
mpsc::Sender<asset_status::Message>,
mpsc::Sender<backfill::Message>,
) {
let guard = Arc::new(RwLock::new(Guard::new()));
let guard = Arc::new(RwLock::new(Guard {
symbols: HashSet::new(),
pending_subscriptions: HashMap::new(),
pending_unsubscriptions: HashMap::new(),
}));
let websocket_url = match thread_type {
ThreadType::Bars(Class::UsEquity) => format!(

26
src/threads/data/websocket.rs
View File

@@ -2,7 +2,7 @@ use super::{backfill, Guard, ThreadType};
use crate::{
config::Config,
database,
types::{alpaca::websocket, Bar, News, Subset},
types::{alpaca::websocket, news::Prediction, Bar, News, Subset},
};
use futures_util::{
stream::{SplitSink, SplitStream},
@@ -19,6 +19,7 @@ use tokio::{
net::TcpStream,
spawn,
sync::{mpsc, Mutex, RwLock},
task::spawn_blocking,
};
use tokio_tungstenite::{tungstenite, MaybeTlsStream, WebSocketStream};
@@ -93,6 +94,7 @@ async fn handle_websocket_message(
}
#[allow(clippy::significant_drop_tightening)]
#[allow(clippy::too_many_lines)]
async fn handle_parsed_websocket_message(
app_config: Arc<Config>,
thread_type: ThreadType,
@@ -195,6 +197,28 @@ async fn handle_parsed_websocket_message(
"{:?} - Received news for {:?}: {}.",
thread_type, news.symbols, news.time_created
);
let app_config_clone = app_config.clone();
let input = format!("{}\n\n{}", news.headline, news.content);
let prediction = spawn_blocking(move || {
app_config_clone
.sequence_classifier
.lock()
.unwrap()
.predict(vec![input.as_str()])
.into_iter()
.map(|label| Prediction::try_from(label).unwrap())
.collect::<Vec<_>>()[0]
})
.await
.unwrap();
let news = News {
sentiment: prediction.sentiment,
confidence: prediction.confidence,
..news
};
database::news::upsert(&app_config.clickhouse_client, &news).await;
}
websocket::incoming::Message::Success(_) => {}

10
src/types/alpaca/api/incoming/asset.rs
View File

@@ -1,7 +1,7 @@
use crate::types::{self, alpaca::api::impl_from_enum};
use serde::{Deserialize, Serialize};
use serde::Deserialize;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Class {
UsEquity,
@@ -10,7 +10,7 @@ pub enum Class {
impl_from_enum!(types::Class, Class, UsEquity, Crypto);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum Exchange {
Amex,
@@ -36,7 +36,7 @@ impl_from_enum!(
Crypto
);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum Status {
Active,
@@ -44,7 +44,7 @@ pub enum Status {
}
#[allow(clippy::struct_excessive_bools)]
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Deserialize)]
pub struct Asset {
pub id: String,
pub class: Class,

6
src/types/alpaca/api/incoming/bar.rs
View File

@@ -1,9 +1,9 @@
use crate::types;
use serde::{Deserialize, Serialize};
use serde::Deserialize;
use std::collections::HashMap;
use time::OffsetDateTime;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Deserialize)]
pub struct Bar {
#[serde(rename = "t")]
#[serde(with = "time::serde::rfc3339")]
@@ -40,7 +40,7 @@ impl From<(Bar, String)> for types::Bar {
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Deserialize)]
pub struct Message {
pub bars: HashMap<String, Vec<Bar>>,
pub next_page_token: Option<String>,

4
src/types/alpaca/api/incoming/clock.rs
View File

@@ -1,7 +1,7 @@
use serde::{Deserialize, Serialize};
use serde::Deserialize;
use time::OffsetDateTime;
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
pub struct Clock {
#[serde(with = "time::serde::rfc3339")]
pub timestamp: OffsetDateTime,

42
src/types/alpaca/api/incoming/news.rs
View File

@@ -1,9 +1,8 @@
use crate::{types, utils::normalize_news_content};
use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use serde::Deserialize;
use time::OffsetDateTime;
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum ImageSize {
Thumb,
@@ -11,14 +10,13 @@ pub enum ImageSize {
Large,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
pub struct Image {
pub size: ImageSize,
pub url: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde_as]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
pub struct News {
pub id: i64,
#[serde(with = "time::serde::rfc3339")]
@@ -28,17 +26,11 @@ pub struct News {
#[serde(rename = "updated_at")]
pub time_updated: OffsetDateTime,
pub symbols: Vec<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub headline: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub author: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub source: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub summary: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub content: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub headline: String,
pub author: String,
pub source: String,
pub summary: String,
pub content: String,
pub url: Option<String>,
pub images: Vec<Image>,
}
@@ -50,17 +42,19 @@ impl From<News> for types::News {
time_created: news.time_created,
time_updated: news.time_updated,
symbols: news.symbols,
headline: normalize_news_content(news.headline),
author: normalize_news_content(news.author),
source: normalize_news_content(news.source),
summary: normalize_news_content(news.summary),
content: normalize_news_content(news.content),
url: news.url,
headline: normalize_news_content(&news.headline),
author: normalize_news_content(&news.author),
source: normalize_news_content(&news.source),
summary: normalize_news_content(&news.summary),
content: normalize_news_content(&news.content),
sentiment: types::news::Sentiment::Neutral,
confidence: 0.0,
url: news.url.unwrap_or_default(),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
pub struct Message {
pub news: Vec<News>,
pub next_page_token: Option<String>,

4
src/types/alpaca/websocket/incoming/bar.rs
View File

@@ -1,8 +1,8 @@
use crate::types;
use serde::{Deserialize, Serialize};
use serde::Deserialize;
use time::OffsetDateTime;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Deserialize)]
pub struct Message {
#[serde(rename = "t")]
#[serde(with = "time::serde::rfc3339")]

4
src/types/alpaca/websocket/incoming/error.rs
View File

@@ -1,6 +1,6 @@
use serde::{Deserialize, Serialize};
use serde::Deserialize;
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Message {
pub code: u16,

4
src/types/alpaca/websocket/incoming/mod.rs
View File

@@ -4,9 +4,9 @@ pub mod news;
pub mod subscription;
pub mod success;
use serde::{Deserialize, Serialize};
use serde::Deserialize;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Deserialize)]
#[serde(tag = "T")]
pub enum Message {
#[serde(rename = "success")]

36
src/types/alpaca/websocket/incoming/news.rs
View File

@@ -1,10 +1,8 @@
use crate::{types, utils::normalize_news_content};
use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use serde::Deserialize;
use time::OffsetDateTime;
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde_as]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
pub struct Message {
pub id: i64,
#[serde(with = "time::serde::rfc3339")]
@@ -14,17 +12,11 @@ pub struct Message {
#[serde(rename = "updated_at")]
pub time_updated: OffsetDateTime,
pub symbols: Vec<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub headline: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub author: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub source: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub summary: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub content: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub headline: String,
pub author: String,
pub source: String,
pub summary: String,
pub content: String,
pub url: Option<String>,
}
@@ -35,12 +27,14 @@ impl From<Message> for types::News {
time_created: news.time_created,
time_updated: news.time_updated,
symbols: news.symbols,
headline: normalize_news_content(news.headline),
author: normalize_news_content(news.author),
source: normalize_news_content(news.source),
summary: normalize_news_content(news.summary),
content: normalize_news_content(news.content),
url: news.url,
headline: normalize_news_content(&news.headline),
author: normalize_news_content(&news.author),
source: normalize_news_content(&news.source),
summary: normalize_news_content(&news.summary),
content: normalize_news_content(&news.content),
sentiment: types::news::Sentiment::Neutral,
confidence: 0.0,
url: news.url.unwrap_or_default(),
}
}
}

8
src/types/alpaca/websocket/incoming/subscription.rs
View File

@@ -1,6 +1,6 @@
use serde::{Deserialize, Serialize};
use serde::Deserialize;
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MarketMessage {
pub trades: Vec<String>,
@@ -14,13 +14,13 @@ pub struct MarketMessage {
pub cancel_errors: Option<Vec<String>>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct NewsMessage {
pub news: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(untagged)]
pub enum Message {
Market(MarketMessage),

4
src/types/alpaca/websocket/incoming/success.rs
View File

@@ -1,6 +1,6 @@
use serde::{Deserialize, Serialize};
use serde::Deserialize;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(tag = "msg")]
#[serde(rename_all = "camelCase")]
pub enum Message {

4
src/types/asset.rs
View File

@@ -4,14 +4,14 @@ use serde_repr::{Deserialize_repr, Serialize_repr};
use time::OffsetDateTime;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
#[repr(u8)]
#[repr(i8)]
pub enum Class {
UsEquity = 1,
Crypto = 2,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
#[repr(u8)]
#[repr(i8)]
pub enum Exchange {
Amex = 1,
Arca = 2,

65
src/types/news.rs
View File

@@ -1,10 +1,49 @@
use clickhouse::Row;
use rust_bert::pipelines::sequence_classification::Label;
use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use serde_repr::{Deserialize_repr, Serialize_repr};
use std::str::FromStr;
use time::OffsetDateTime;
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Row)]
#[serde_as]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
#[repr(i8)]
pub enum Sentiment {
Positive = 1,
Neutral = 0,
Negative = -1,
}
impl FromStr for Sentiment {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"positive" => Ok(Self::Positive),
"neutral" => Ok(Self::Neutral),
"negative" => Ok(Self::Negative),
_ => Err(()),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Prediction {
pub sentiment: Sentiment,
pub confidence: f64,
}
impl TryFrom<Label> for Prediction {
type Error = ();
fn try_from(label: Label) -> Result<Self, Self::Error> {
Ok(Self {
sentiment: Sentiment::from_str(&label.text)?,
confidence: label.score,
})
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Row)]
pub struct News {
pub id: i64,
#[serde(with = "clickhouse::serde::time::datetime")]
@@ -12,16 +51,12 @@ pub struct News {
#[serde(with = "clickhouse::serde::time::datetime")]
pub time_updated: OffsetDateTime,
pub symbols: Vec<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub headline: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub author: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub source: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub summary: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub content: Option<String>,
#[serde_as(as = "NoneAsEmptyString")]
pub url: Option<String>,
pub headline: String,
pub author: String,
pub source: String,
pub summary: String,
pub content: String,
pub sentiment: Sentiment,
pub confidence: f64,
pub url: String,
}

11
src/utils/news.rs
View File

@@ -1,10 +1,7 @@
use html_escape::decode_html_entities;
use regex::Regex;
pub fn normalize_news_content(content: Option<String>) -> Option<String> {
content.as_ref()?;
let content = content.unwrap();
pub fn normalize_news_content(content: &str) -> String {
let re_tags = Regex::new("<[^>]+>").unwrap();
let re_spaces = Regex::new("[\\u00A0\\s]+").unwrap();
@@ -14,9 +11,5 @@ pub fn normalize_news_content(content: Option<String>) -> Option<String> {
let content = decode_html_entities(&content);
let content = content.trim();
if content.is_empty() {
None
} else {
Some(content.to_string())
}
content.to_string()
}