!pip install transformers torch datasets vaderSentiment pandas matplotlib seaborn tqdm --quiet

import torch
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer
from tqdm.auto import tqdm
import time
import re

%matplotlib inline
plt.rcParams['figure.figsize'] = (12, 5)
plt.rcParams['axes.spines.top'] = False
plt.rcParams['axes.spines.right'] = False

DEVICE = 0 if torch.cuda.is_available() else -1
DEVICE_NAME = 'GPU' if DEVICE == 0 else 'CPU'
print(f'Running on: {DEVICE_NAME}')
if DEVICE == 0:
    print(f'GPU: {torch.cuda.get_device_name(0)}')

# ─── MODEL CONFIGURATION ────────────────────────────────────────────────────
PRIMARY_MODEL   = 'cardiffnlp/twitter-roberta-base-sentiment-latest'
FAST_MODEL      = 'distilbert-base-uncased-finetuned-sst-2-english'
BATCH_SIZE      = 32    # Reduce to 16 if you hit CUDA OOM errors
MAX_TOKEN_LEN   = 128   # Max tokens per text (truncate longer texts)
# ────────────────────────────────────────────────────────────────────────────


def load_sentiment_pipeline(model_name: str, device: int = -1) -> pipeline:
    """
    Load a HuggingFace text-classification pipeline.

    The pipeline handles tokenization and inference in one call.
    Downloads model weights on first use (~500MB for RoBERTa).

    Parameters
    ----------
    model_name : HuggingFace model identifier
    device     : 0 for GPU, -1 for CPU

    Returns
    -------
    transformers.Pipeline object
    """
    print(f'Loading model: {model_name} on {"GPU" if device == 0 else "CPU"}...')
    t0 = time.time()
    nlp = pipeline(
        task='text-classification',
        model=model_name,
        device=device,
        truncation=True,
        max_length=MAX_TOKEN_LEN,
        return_all_scores=True
    )
    print(f'Model loaded in {time.time() - t0:.1f}s')
    return nlp


bert_pipeline = load_sentiment_pipeline(PRIMARY_MODEL, DEVICE)

SAMPLE_TEXTS = [
    # Strongly positive
    'Bitcoin just broke ATH! This bull run is just getting started, we are going to $200k!',
    'Ethereum 2.0 staking yields are incredible. Best passive income in crypto right now.',
    'Just bought more BTC on this dip. Thank you bears for the discount 🙏',
    'Institutional adoption is accelerating. BlackRock ETF approval is massive for the space.',
    'DeFi TVL hit a new record. This ecosystem is unstoppable.',
    # Moderately positive
    'Crypto markets are looking stable today, slight upward trend across the board.',
    'Bitcoin dominance rising — historically a good sign before altseason.',
    'Layer 2 solutions are finally delivering on their promises. Fees are much lower.',
    # Neutral / informational
    'Bitcoin volume was below average today at 24 billion USD.',
    'The Federal Reserve meets next week. Traders watching for rate decision signals.',
    'Ethereum gas fees averaged 15 gwei over the past 24 hours.',
    'BTC is currently trading at $65,400, down 0.3% from yesterday.',
    # Moderately negative
    'Crypto markets pulled back today after regulatory news from the SEC.',
    'Bitcoin struggling to hold support at $60k. Bears testing the level again.',
    'The altcoin market looks weak. Most coins are bleeding vs BTC.',
    # Strongly negative
    'CRASH! BTC down 15% in one hour. Exchange liquidations everywhere. This is a disaster.',
    'Another exchange hack. $300M stolen. This is why I dont trust centralized platforms.',
    'The crypto market is dead. Retail is gone. Only whales manipulating price now.',
    'SEC is destroying crypto innovation. These regulations will kill the entire industry.',
    'Absolute bloodbath today. Lost 40% of my portfolio. I am done with crypto forever.',
    # Ambiguous / sarcastic
    'Oh great, another 10% down day. Totally unexpected 🙄',
    'Sure, just buy the dip they said. Great advice.',
    'Bitcoin is a hedge against inflation... at -60% YTD 😂'
]

sample_df = pd.DataFrame({'text': SAMPLE_TEXTS})
print(f'Sample corpus: {len(sample_df)} texts')
sample_df.head(5)

def preprocess_for_bert(text: str, max_chars: int = 512) -> str:
    """
    Minimal preprocessing before BERT tokenization.

    BERT handles raw text well — we only remove URLs and excessive whitespace.
    Preserving punctuation and capitalization is important for context.

    Parameters
    ----------
    text     : Raw text string
    max_chars: Truncate to this many characters before tokenization

    Returns
    -------
    Lightly cleaned text string
    """
    if not isinstance(text, str):
        return ''
    text = re.sub(r'http\S+|www\.\S+', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text[:max_chars]


def run_bert_batch_inference(
    texts: list[str],
    nlp_pipeline,
    batch_size: int = 32
) -> pd.DataFrame:
    """
    Run batch sentiment inference using a HuggingFace pipeline.

    Returns a DataFrame with one row per input text and columns
    for each label's confidence score plus the winning label.

    Parameters
    ----------
    texts        : List of text strings to classify
    nlp_pipeline : Loaded HuggingFace pipeline
    batch_size   : Number of texts per inference batch

    Returns
    -------
    DataFrame with columns per label score + bert_label + bert_confidence
    """
    cleaned = [preprocess_for_bert(t) for t in texts]

    t0 = time.time()
    all_results = []
    for i in tqdm(range(0, len(cleaned), batch_size), desc='BERT inference'):
        batch = cleaned[i : i + batch_size]
        outputs = nlp_pipeline(batch, batch_size=len(batch))
        # The pipeline returns List[List[Dict]] when return_all_scores=True and batching.
        # We need to append each inner list (scores for one text) individually.
        for text_output in outputs:
            all_results.append(text_output)
    elapsed = time.time() - t0

    print(f'Processed {len(texts)} texts in {elapsed:.2f}s ({len(texts)/elapsed:.0f} texts/sec)')

    rows = []
    for result in all_results:
        # result is now correctly List[Dict], e.g., [{'label': 'LABEL_0', 'score': 0.9}, {'label': 'LABEL_1', 'score': 0.1}]
        row = {item['label'].lower().replace('label_', ''): item['score'] for item in result}
        top = max(result, key=lambda x: x['score'])
        row['bert_label']      = top['label'].lower().replace('label_', '')
        row['bert_confidence'] = top['score']
        rows.append(row)

    return pd.DataFrame(rows)

import pandas as pd
from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer
from tqdm.auto import tqdm
import time
import re

vader = SentimentIntensityAnalyzer()

def preprocess_for_bert(text: str, max_chars: int = 512) -> str:
    if not isinstance(text, str):
        return ''
    text = re.sub(r'http\S+|www\.\S+', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text[:max_chars]

def run_bert_batch_inference(
    texts: list[str],
    nlp_pipeline,
    batch_size: int = 32
) -> pd.DataFrame:
    cleaned = [preprocess_for_bert(t) for t in texts]
    all_results = []

    t0 = time.time()
    for i in tqdm(range(0, len(cleaned), batch_size), desc='BERT inference'):
        batch = cleaned[i : i + batch_size]
        outputs = nlp_pipeline(batch, batch_size=len(batch))

        if not outputs: continue

        # Standardize output to List[List[Dict]]
        if isinstance(outputs[0], dict):
            # Case: Pipeline returned List[Dict] (one dict per text)
            for item in outputs:
                all_results.append([item])
        elif isinstance(outputs[0], list):
            # Case: Pipeline returned List[List[Dict]] (multiple scores per text)
            all_results.extend(outputs)

    elapsed = time.time() - t0
    print(f'Processed {len(texts)} texts in {elapsed:.2f}s')

    rows = []
    for result in all_results:
        # Extract all available scores in the list
        row = {item['label'].lower().replace('label_', ''): item['score'] for item in result}
        # Determine top label
        top = max(result, key=lambda x: x['score'])
        row['bert_label'] = top['label'].lower().replace('label_', '')
        row['bert_confidence'] = top['score']
        rows.append(row)

    return pd.DataFrame(rows)

def add_vader_comparison(df: pd.DataFrame, text_col: str = 'text') -> pd.DataFrame:
    df = df.copy()
    df['vader_compound'] = df[text_col].apply(lambda t: vader.polarity_scores(str(t))['compound'])
    df['vader_label'] = df['vader_compound'].apply(
        lambda c: 'positive' if c >= 0.05 else ('negative' if c <= -0.05 else 'neutral')
    )
    df['models_agree'] = df['bert_label'] == df['vader_label']
    print(f"BERT vs VADER agreement: {df['models_agree'].mean() * 100:.1f}%")
    return df

# Execution
bert_scores = run_bert_batch_inference(sample_df['text'].tolist(), bert_pipeline, BATCH_SIZE)
result_df = pd.concat([sample_df.reset_index(drop=True), bert_scores.reset_index(drop=True)], axis=1)
compared_df = add_vader_comparison(result_df)

print(f'\nDisagreements ({len(compared_df[~compared_df["models_agree"]])} cases):')
compared_df[~compared_df['models_agree']][['text', 'bert_label', 'bert_confidence', 'vader_label', 'vader_compound']].head(10)

def plot_bert_vader_comparison(df: pd.DataFrame) -> None:
    """
    Scatter plot comparing BERT confidence vs VADER compound score.
    Points are colored by BERT label.

    Parameters
    ----------
    df : DataFrame with both BERT and VADER scores
    """
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Left: Scatter — VADER compound vs BERT (positive conf - negative conf)
    bert_score = df.get('positive', pd.Series(0, index=df.index)) - df.get('negative', pd.Series(0, index=df.index))
    colors = df['bert_label'].map({'positive': '#2ecc71', 'neutral': '#95a5a6', 'negative': '#e74c3c'})
    axes[0].scatter(df['vader_compound'], bert_score, c=colors, s=80, alpha=0.8, edgecolors='white')
    axes[0].axhline(0, color='gray', linestyle='--', linewidth=0.8)
    axes[0].axvline(0, color='gray', linestyle='--', linewidth=0.8)
    axes[0].set_xlabel('VADER Compound Score')
    axes[0].set_ylabel('BERT Score (pos − neg confidence)')
    axes[0].set_title('VADER vs BERT Agreement\n(points far from diagonal = disagreements)', fontsize=11, fontweight='bold')

    # Right: Confusion matrix style label comparison
    cross = pd.crosstab(df['vader_label'], df['bert_label'],
                         rownames=['VADER'], colnames=['BERT'])
    sns.heatmap(cross, annot=True, fmt='d', cmap='Blues', ax=axes[1],
                linewidths=0.5, cbar=False)
    axes[1].set_title('Label Agreement Heatmap', fontsize=11, fontweight='bold')

    plt.suptitle('BERT vs VADER Sentiment Comparison', fontsize=13)
    plt.tight_layout()
    plt.show()


plot_bert_vader_comparison(compared_df)

def run_zero_shot_classification(
    texts: list[str],
    candidate_labels: list[str],
    model_name: str = 'facebook/bart-large-mnli',
    device: int = -1
) -> pd.DataFrame:
    """
    Classify texts against custom labels using zero-shot NLI.

    No labeled training data is needed — labels are provided as strings.
    Returns a score for each label for each input text.

    Parameters
    ----------
    texts             : List of text strings to classify
    candidate_labels  : List of category strings (e.g., ['bullish', 'bearish', 'FUD'])
    model_name        : Zero-shot NLI model from HuggingFace
    device            : 0 for GPU, -1 for CPU

    Returns
    -------
    DataFrame with one row per text and columns for each label's score
    """
    zsc = pipeline('zero-shot-classification', model=model_name, device=device)

    rows = []
    for text in tqdm(texts, desc='Zero-shot classification'):
        result = zsc(text[:400], candidate_labels=candidate_labels)
        row = dict(zip(result['labels'], result['scores']))
        row['top_label'] = result['labels'][0]
        rows.append(row)

    return pd.DataFrame(rows)


CRYPTO_LABELS = ['bullish price prediction', 'bearish price prediction',
                 'regulatory risk', 'technical analysis', 'market news']

# Use a small subset to save quota
zsc_results = run_zero_shot_classification(
    SAMPLE_TEXTS[:10], CRYPTO_LABELS, device=DEVICE
)
pd.concat([pd.Series(SAMPLE_TEXTS[:10], name='text'), zsc_results], axis=1)

def benchmark_models(
    texts: list[str],
    models: dict,
    device: int = -1
) -> pd.DataFrame:
    """
    Benchmark inference speed of multiple HuggingFace sentiment models.

    Parameters
    ----------
    texts  : List of text strings to use as benchmark input
    models : Dictionary of {model_name: model_display_name}
    device : 0 for GPU, -1 for CPU

    Returns
    -------
    DataFrame with timing results per model
    """
    results = []
    for model_id, display_name in models.items():
        print(f'Benchmarking: {display_name}...')
        nlp = load_sentiment_pipeline(model_id, device)
        t0  = time.time()
        _   = nlp(texts, batch_size=32)
        elapsed = time.time() - t0
        results.append({
            'model':         display_name,
            'texts_processed': len(texts),
            'total_seconds': round(elapsed, 2),
            'texts_per_sec': round(len(texts) / elapsed)
        })
        del nlp
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    return pd.DataFrame(results)


MODELS_TO_BENCHMARK = {
    'distilbert-base-uncased-finetuned-sst-2-english': 'DistilBERT (fast)',
    'cardiffnlp/twitter-roberta-base-sentiment-latest': 'Twitter-RoBERTa'
}

benchmark_texts = (SAMPLE_TEXTS * 5)[:100]
bench_df = benchmark_models(benchmark_texts, MODELS_TO_BENCHMARK, DEVICE)
bench_df

compared_df.to_csv('bert_sentiment_results.csv', index=False)
print(f'Results saved: bert_sentiment_results.csv ({len(compared_df)} rows)')
print(f'\nColumns: {list(compared_df.columns)}')