Update finetuned_model.py

finetuned_model.py

@@ -78,7 +78,7 @@ for _, row in df_yearly.iterrows():
         "summary": f"In {year}, the S&P 500 averaged {sp500:.2f} with a {return_val:.1f}% annual return and a {real_return:.1f}% real return."
     })
 
-# Period-specific questions (1-year, 3-year, 5-year, 10-year, and recent ranges)
+# Period-specific questions (1-year, 3-year, 5-year, 10-year, and custom ranges)
 years = df_yearly['Year'].unique()
 for year in years:
     for duration in [1, 3, 5, 10]:
@@ -97,8 +97,8 @@ for year in years:
                 "answer": f"The S&P 500’s {duration}-year average annual inflation-adjusted return from {start_year} to {end_year} was approximately {avg_real_return:.1f}%."
             })
 
-# Custom period questions, including recent periods
-custom_periods = [(2000, 2010), (2011, 2016), (2010, 2020), (2000, 2008), (2015, 2024), (2020, 2022), (2020, 2024)]
+# Custom period questions
+custom_periods = [(2000, 2010), (2011, 2016), (2010, 2020), (2000, 2008), (2015, 2024)]
 for start_year, end_year in custom_periods:
     df_period = df_yearly[(df_yearly['Year'] >= start_year) & (df_yearly['Year'] <= end_year)]
     if not df_period.empty:
@@ -139,12 +139,8 @@ qa_pairs.append({
     "answer": "The S&P 500 index fund’s average annual return is approximately 10–12% over the long term (1927–2025), including dividends, based on historical data."
 })
 qa_pairs.append({
-    "question": "What was the average annual return of the S&P 500 between 2020 and 2022?",
-    "answer": "The S&P 500’s average annual return from 2020 to 2022 was approximately 8.3%, including dividends, with significant volatility due to the COVID-19 recovery and 2022 bear market."
-})
-qa_pairs.append({
-    "question": "What was the average annual return of the S&P 500 in the past 5 years?",
-    "answer": "The S&P 500’s average annual return from 2020 to 2024 was approximately 10.5%, including dividends, based on historical data."
+    "question": "What was the average annual return of the S&P 500 between 2010 and 2020?",
+    "answer": "The S&P 500’s average annual return from 2010 to 2020 was approximately 13.6%, including dividends, driven by post-financial crisis recovery."
 })
 
 # Save to JSON
@@ -181,8 +177,8 @@ training_args = TrainingArguments(
     output_dir="./finetuned_model",
     evaluation_strategy="epoch",
     learning_rate=1e-5,
-    per_device_train_batch_size=16,
-    per_device_eval_batch_size=16,
+    per_device_train_batch_size=8,  # Increased for faster training
+    per_device_eval_batch_size=8,
     num_train_epochs=7,
     weight_decay=0.01,
     logging_steps=10,
@@ -210,15 +206,7 @@ trainer.save_model("./finetuned_model")
 tokenizer.save_pretrained("./finetuned_model")
 
 # Test the model
-input_text = "What was the average annual return of the S&P 500 between 2020 and 2022?"
-inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-inputs = {k: v.to("cpu") for k, v in inputs.items()}
-with torch.inference_mode():
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=20,
-        repetition_penalty=3.0,
-        no_repeat_ngram_size=2,
-        pad_token_id=tokenizer.eos_token_id
-    )
+input_text = "What was the average annual return of the S&P 500 between 2010 and 2020?"
+inputs = tokenizer(input_text, return_tensors="pt")
+outputs = model.generate(**inputs, max_new_tokens=40)
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))