import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
from sklearn.preprocessing import StandardScaler

class RecommendModel(nn.Module):
    def __init__(self, input_size, embedding_size):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_size, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, embedding_size)
        )
        
    def forward(self, x):
        return self.encoder(x)

def train_model():
    # 加载数据
    videos_df = pd.read_csv('videos.csv')
    behaviors_df = pd.read_csv('behaviors.csv')
    
    # 数据预处理
    scaler = StandardScaler()
    features = scaler.fit_transform(videos_df[['duration', 'view_count', 'like_count']])
    
    # 创建模型
    model = RecommendModel(input_size=features.shape[1], embedding_size=128)
    criterion = nn.CosineEmbeddingLoss()
    optimizer = optim.Adam(model.parameters())
    
    # 训练模型
    for epoch in range(100):
        # ... 训练代码
        
    # 保存模型
    torch.save(model.state_dict(), '../src/main/resources/models/recommend_model.pt')

if __name__ == '__main__':
    train_model()