Self-Learning LLMs for Stock Forecasting: A Python Implementation with Direct Preference Optimization

def get_data():
    try:
         # Load dataset from Hugging Face
        dataset = load_dataset("oliverwang15/us_stock_news_with_price")
        df = dataset["train"].to_pandas()
        print("Available Columns:", df.columns)

        # Extract necessary columns and drop missing values
        df = df[["stock", "title", "content", "ts_0", "ts_-1", "trading_date", "exact_trading_date"]].dropna()

        # Calculate news effect: positive if next_day_price >= news_day_price, else negative
        df["news_effect"] = df.apply(
            lambda row: "positive" if row["ts_0"] >= row["ts_-1"] else "negative", axis=1
        )

        # Rename columns for clarity
        df_news = df.rename(
            {
                "stock": "ticker",
                "title": "news_title",
                "content": "news_summary",
                "ts_0": "next_day_price",
                "ts_-1": "news_day_price",
            },
            axis=1,
        )
        print("Final Columns:", df_news.columns)
        return df_news
    except Exception as e:
        print(f"Error loading or processing data: {e}")
        return None

def create_database():
    # Create SQLite database connection
    conn = sqlite3.connect("stock_news.db")
    cursor = conn.cursor()

    # Enable foreign key support
    cursor.execute("PRAGMA foreign_keys = ON;")

    # Create tables in SQLite
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS stock_news (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            ticker TEXT,
            news_title TEXT,
            news_summary TEXT,
            next_day_price FLOAT,
            news_day_price FLOAT,
            trading_date TEXT,
            exact_trading_date TEXT,
            news_effect TEXT
        );
    """
    )

    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS news_sentiment (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            news_id INTEGER,
            sentiment TEXT,
            news_explanation TEXT,
            FOREIGN KEY (news_id) REFERENCES stock_news(id) ON DELETE CASCADE
        );
    """
    )

    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS news_forecast (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            news_id INTEGER,
            ticker TEXT,
            forecast_1 TEXT,
            probability_1 FLOAT,
            forecast_2 TEXT,
            probability_2 FLOAT,
            best_forecast TEXT,
            FOREIGN KEY (news_id) REFERENCES news_sentiment(news_id) ON DELETE CASCADE
        );
    """
    )
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS model_updates (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT,
            update_date TEXT,
            brier_score FLOAT
        );
    """)

    ## insert the hugging face data
    def insert_data(df_data):
        conn = sqlite3.connect("stock_news.db")
        cursor = conn.cursor()
        for _, row in df_data.iterrows():
            cursor.execute("""
                INSERT INTO stock_news (ticker, news_title, news_summary, next_day_price, news_day_price, news_effect)
                VALUES (?, ?, ?, ?, ?, ?);
            """, (row["ticker"], row["news_title"], row["news_summary"], row["next_day_price"],
                    row["news_day_price"], row["news_effect"]))
        conn.commit()
        conn.close()

def analyze_sentiment_and_explain(news_title, news_summary):
    """Use the model to analyze sentiment and provide an explanation in JSON format."""
    prompt = f"""
    Analyze the sentiment of the following stock news article and explain your analysis

    Title: {news_title}
    Summary: {news_summary}

    Provide the response in **valid JSON format**:
    {{
        "sentiment": "<positive/negative/neutral>",
        "explanation": "<Brief explanation>"
    }}
    """
    response: ChatResponse = chat(model='qwen2.5', messages=[
        {
            'role': 'user',
            'content': f'{prompt}'
        }])


    response_text = response['message']['content'].strip()

    try:
        sentiment_data = json.loads(response_text)  # Parse JSON response
        sentiment = sentiment_data.get(
            "sentiment", "neutral"
        )  # Default to Neutral if sentiment is not provided
        sentiment = sentiment.lower()
        explanation = sentiment_data.get("explanation", "No explanation available.")
    except json.JSONDecodeError as e:
        # If JSON parsing fails, default to Neutral sentiment
        sentiment, explanation = "neutral", f"Sentiment could not be determined. {str(e)}"

    return sentiment, explanation

def get_sentiment():
    conn = sqlite3.connect("stock_news.db")
    cursor = conn.cursor()
    cursor.execute("""SELECT id, ticker, news_title, news_summary
                        FROM stock_news
                        LIMIT 20;""") # just a subset of the data for now
    news_data = cursor.fetchall()

    for news_id, ticker, title, summary in news_data:
        # call teh model with our data
        sentiment, explanation = analyze_sentiment_and_explain(title, summary)

        # Insert sentiment & explanation into news_sentiment table
        cursor.execute(
            """
            INSERT INTO news_sentiment (news_id, sentiment, news_explanation)
            VALUES (?, ?, ?);
        """,
            (news_id, sentiment, explanation),
        )
    conn.commit()
    conn.close()

def extract_json_from_tags(text):
    """Extract JSON content between ```json and ``` tags."""
    match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL)
    return match.group(1) if match else None


def forecast_sentiment(ticker, sentiment, news_explanation):
    """Prompt GPT-4 to generate a stock movement forecast as 'positive' or 'negative'."""
    prompt = f"""
    Given the question: Is this forecast {sentiment} for stock {ticker}.
    And the following news summary about the stock: {news_explanation}

    Provide a JSON response with the following:
    {{
        "forecast": "<positive/negative>",
        "probability": <A number between 0 and 1>
        "explanation": "<Brief explanation>"
    }}
    """



    response: ChatResponse = chat(model='qwen2.5', messages=[
        {
            'role': 'user',
            'content': f'{prompt}'
        }])

    response_text = response['message']['content'].strip()
    if "```json" in response_text:
        response_text = extract_json_from_tags(response_text)

    try:
        forecast_data = json.loads(response_text)  # Parse JSON response
        forecast = forecast_data.get(
            "forecast", "missing"
        )  # Default to "missing" if missing
        probability = forecast_data.get("probability", 0.0)
    except json.JSONDecodeError:
        forecast, probability = "missing", 0.0

    return forecast, probability

    # Fetch sentiment data for forecasting along with the stock ticker


def gen_forecast():
    conn = sqlite3.connect("stock_news.db")
    cursor = conn.cursor()
    cursor.execute(
        """
        SELECT news_sentiment.news_id,
        stock_news.ticker, 
        news_sentiment.sentiment, 
        news_sentiment.news_explanation 
        FROM news_sentiment
        JOIN stock_news ON news_sentiment.news_id = stock_news.id;
    """
    )
    sentiment_data = cursor.fetchall()
    for news_id, ticker, sentiment, explanation in sentiment_data:
        print(f"Generating forecasts for id: {news_id} stock {ticker} with sentiment {sentiment}...")
        forecast_1, probability_1 = forecast_sentiment(
            ticker, "positive", explanation
        )
        forecast_2, probability_2 = forecast_sentiment(
            ticker, "negative", explanation
        )

        if probability_1 > probability_2:
            best_forecast = forecast_1
        else:
            best_forecast = forecast_2

        # Insert forecasts into news_forecast table
        cursor.execute(
            """
            INSERT INTO news_forecast (news_id, ticker, forecast_1, probability_1, forecast_2, probability_2, best_forecast)
            VALUES (?, ?, ?, ?, ?, ?, ?);
        """,
            (news_id, ticker, forecast_1, probability_1, forecast_2, probability_2, best_forecast),
        )

    # Commit forecasts to database
    conn.commit()
    conn.close()

def ranking_metric(prediction, actual_outcome):
    """Calculate the absolute error between prediction and actual outcome."""
    return abs(prediction - actual_outcome)

# Convert forecast strings to float probabilities
df["forecast_1"] = df["forecast_1"].str.extract(r'(\d\.\d+)').astype(float)
df["forecast_2"] = df["forecast_2"].str.extract(r'(\d\.\d+)').astype(float)

# Compute ranking scores
df["rank_1"] = df.apply(lambda row: ranking_metric(row["forecast_1"], row["outcome"]), axis=1)
df["rank_2"] = df.apply(lambda row: ranking_metric(row["forecast_2"], row["outcome"]), axis=1)

# Determine the better forecast
df["better_forecast"] = df.apply(lambda row: "Forecast 1" if row["rank_1"] < row["rank_2"] else "Forecast 2", axis=1)

# Display dataset with rankings
tools.display_dataframe_to_user(name="Ranked Forecasts", dataframe=df)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import get_peft_model, LoraConfig, TaskType

def fine_tune_dpo(df):
    model_name = "mistralai/Mistral-7B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
    
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=["q_proj", "v_proj"]
    )
    model = get_peft_model(model, peft_config)
    
    train_texts = [f"News: {row['content']}\nSentiment: {row['news_effect']}\n" for _, row in df.iterrows()]
    train_labels = [row['news_effect'] for _, row in df.iterrows()]
    
    tokenized_data = tokenizer(train_texts, padding=True, truncation=True, return_tensors="pt")
    
    training_args = TrainingArguments(
        output_dir="./dpo_model",
        per_device_train_batch_size=4,
        num_train_epochs=3,
        logging_dir="./logs",
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_data,
    )
    trainer.train()
    model.save_pretrained("./dpo_model")
    tokenizer.save_pretrained("./dpo_model")

# Train the model
fine_tune_dpo(df)

def brier_score(predictions, outcomes):
    """Compute the Brier Score for probabilistic forecasts."""
    return ((predictions - outcomes) ** 2).mean()

brier = brier_score(df["forecast_1"], df["outcome"])
print(f"Brier Score: {brier:.4f}")