ProfilingAI / src /core /data_generator.py
Sandrine Guétin
Version propre de DeepVest
2106f78
raw
history blame
2.22 kB
import pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM
class ProfileGenerator:
def __init__(self):
# Charger Open R1 pour la partie profiling
self.tokenizer = AutoTokenizer.from_pretrained("open-r1")
self.model = AutoModelForCausalLM.from_pretrained("open-r1")
# Configuration des profils
self.risk_levels = ["Prudent", "Modéré", "Dynamique"]
self.goals = ["Retraite", "Immobilier", "Études", "Épargne"]
def generate_profile_data(self, num_samples=1000):
data = []
for _ in range(num_samples):
# Générer profil de base
profile = {
"age": np.random.randint(18, 70),
"income": np.random.randint(2000, 20000),
"savings": np.random.randint(500, 5000),
"goals": np.random.choice(self.goals, size=2, replace=False),
"risk_tolerance": np.random.choice(self.risk_levels)
}
# Générer description avec Open R1
profile["description"] = self._generate_profile_description(profile)
data.append(profile)
return pd.DataFrame(data)
def _generate_profile_description(self, profile):
# Préparation du prompt pour Open R1
prompt = f"""
Profil investisseur :
- Age: {profile['age']} ans
- Revenu: {profile['income']}€/mois
- Épargne: {profile['savings']}€/mois
- Objectifs: {', '.join(profile['goals'])}
- Tolérance au risque: {profile['risk_tolerance']}
Décrivez ce profil d'investisseur et suggérez une approche personnalisée.
"""
# Générer la description avec Open R1
inputs = self.tokenizer(prompt, return_tensors="pt")
outputs = self.model.generate(**inputs, max_length=200)
return self.tokenizer.decode(outputs[0])
def main():
# Générer données
generator = ProfileGenerator()
df = generator.generate_profile_data()
# Sauvegarder
df.to_csv("src/data/investment_profiles.csv", index=False)
print("Données générées et sauvegardées")
if __name__ == "__main__":
main()