Update finetuned_model.py

finetuned_model.py

@@ -68,21 +68,46 @@ for _, row in df.iterrows():
         "summary": f"On {date}, the S&P 500 closed at {sp500:.2f} with a {return_val:.1f}% annual return and a {real_return:.1f}% real return."
     })
 
-# Add period-specific questions
-# Example periods: 2000–2010, 2011–2016, 2010–2020
-periods = [(2000, 2010), (2011, 2016), (2010, 2020)]
+# Period-specific questions
+periods = [(2000, 2010), (2011, 2016), (2010, 2020), (2018, 2020), (2020, 2023)]
 for start_year, end_year in periods:
     df_period = df[(df['Date'].dt.year >= start_year) & (df['Date'].dt.year <= end_year)]
-    avg_return = df_period['Return'].mean()
-    avg_real_return = df_period['Real Return'].mean()
-    qa_pairs.append({
-        "question": f"What is the average annual growth rate of the S&P 500 from {start_year} to {end_year}?",
-        "answer": f"The S&P 500’s average annual growth rate from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
-    })
-    qa_pairs.append({
-        "question": f"What was the S&P 500’s real return from {start_year} to {end_year}?",
-        "answer": f"The S&P 500’s average annual inflation-adjusted return from {start_year} to {end_year} was approximately {avg_real_return:.1f}%."
-    })
+    if not df_period.empty:
+        avg_return = df_period['Return'].mean()
+        avg_real_return = df_period['Real Return'].mean()
+        qa_pairs.append({
+            "question": f"What is the average annual growth rate of the S&P 500 from {start_year} to {end_year}?",
+            "answer": f"The S&P 500’s average annual growth rate from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
+        })
+        qa_pairs.append({
+            "question": f"What was the S&P 500’s real return from {start_year} to {end_year}?",
+            "answer": f"The S&P 500’s average annual inflation-adjusted return from {start_year} to {end_year} was approximately {avg_real_return:.1f}%."
+        })
+
+# Relative period questions (e.g., last 1, 3, 5 years)
+current_year = 2025
+relative_periods = [1, 3, 5]
+for years in relative_periods:
+    start_year = current_year - years
+    df_period = df[(df['Date'].dt.year >= start_year) & (df['Date'].dt.year <= current_year)]
+    if not df_period.empty:
+        avg_return = df_period['Return'].mean()
+        qa_pairs.append({
+            "question": f"What is the S&P 500 growth rate for the last {years} years?",
+            "answer": f"The S&P 500’s average annual growth rate for the last {years} years (from {start_year} to {current_year}) was approximately {avg_return:.1f}%, including dividends."
+        })
+
+# Investment return questions
+amounts = [1000, 5000, 10000]
+years = [1, 3, 5, 10, 20]
+avg_annual_return = 10.0  # Historical S&P 500 average
+for amount in amounts:
+    for n in years:
+        future_value = amount * (1 + avg_annual_return / 100) ** n
+        qa_pairs.append({
+            "question": f"What will ${amount} be worth in {n} years if invested in the S&P 500?",
+            "answer": f"Assuming a 10% average annual return, ${amount:,.0f} invested in the S&P 500 would grow to approximately ${future_value:,.0f} in {n} years with annual compounding."
+        })
 
 # Add general S&P 500 growth rate question
 qa_pairs.append({
@@ -153,8 +178,7 @@ trainer.save_model("./finetuned_model")
 tokenizer.save_pretrained("./finetuned_model")
 
 # Test the model
-input_text = "What is the average annual growth rate of the S&P 500 from 2000 to 2010?"
+input_text = "What is the S&P 500 growth rate for the last 3 years?"
 inputs = tokenizer(input_text, return_tensors="pt")
 outputs = model.generate(**inputs, max_new_tokens=50)
-
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))