Views
No views yet
rlhf_type=grpo) via ms-swift. During training, each input is augmented with group memory retrieval (member listening history + common group artists).memory/plungin.py):| Reward function | Weight | Description |
|---|---|---|
format_reward | 0.2 | Checks that the <think> <memory> <reasoning> <rec> tags are complete, properly closed, and in the correct order |
recommendation_reward | 0.5 | Computes Hit@k / NDCG@k between the recommendation list and the ground truth |
reasoning_reward | 0.3 | Calls the DeepSeek API to score the quality of the reasoning process |
/root/autodl-tmp) and require a GPU with at least ~16 GB VRAM (Llama-3-8B in fp16).1# Install ms-swift (training framework) and dependencies
2pip install 'ms-swift[llm]' -U
3pip install pandas datasets
4
5# Create the project directory and copy the code
6mkdir -p /root/autodl-tmp/grpo_memory
7# Copy the memory/ folder from this repo (config.py, preprocess_dataset.py,
8# memory_retriever.py, plungin.py, prompt.txt, set_paths.sh, run_grpo_enhanced.sh, csv/)BASE_PATH in memory/config.py and memory/set_paths.sh — all other paths are computed from it automatically:1# config.py
2BASE_PATH = "/root/autodl-tmp" # <- change to your server path1[
2 {
3 "prompt": "Based on the group's listening history and individual music preferences, analyze group and user preference features ... and recommend the top 10 suitable artists from the candidate list with ranking and reasons.",
4 "input": "Group info Group ID:group_common0000 Group size:2 Common artist count:7 Current common artists:Rihanna(pop,rnb,dance),MariahCarey(rnb,pop,female vocalists)...Member listening preferences Member 503:BritneySpears(tags:pop,dance,female vocalists,preference strength:9.6/10)...Candidate artist list 1.HironosukeSatou(random,unknown) 2.TheKillers(indie,indie rock,rock)...",
5 "output": "...reference answer...",
6 "ground_truth": "BritneySpears"
7 }
8]prompt: task instructioninput: contains the group ID, member listening preferences, and the candidate artist list; the memory retriever extracts group_xxx and member IDs from hereoutput / ground_truth: reference answer and the ground-truth artist(s), used by the reward functions to compute hit rates| File | Columns | Purpose |
|---|---|---|
user_artists.csv | userID, artistID, weight | user-artist listening relations (weight = listening weight) |
artists.csv | id, name, url | artist ID → name mapping |
tags.csv | userID, artistID, tagID, day, month, year | tagging records |
group_common0000) and member IDs from each sample's input;memory_retriever.py to retrieve each member's top-3 listening history and the group's top-3 common artists;【记忆检索结果】 (memory retrieval result) block to the original input to build the enhanced dataset;group_id, has_memory, etc.) and print enhancement statistics.1cd /root/autodl-tmp/grpo_memory
2python preprocess_dataset.py
3# Writes the enhanced dataset to the ENHANCED_DATASET path configured in config.py...original input content...
【记忆检索结果】
【User listening history】
User 503: BritneySpears, GleeCast, Keane
User 145: JimSturgess, EllieGoulding, DavidArchuleta
【Group common artists】
Rihanna (average weight: 9)
Please make music recommendations based on the memory information above.1cd /root/autodl-tmp/grpo_memory
2source set_paths.sh
3bash run_grpo_enhanced.sh1swift rlhf \
2 --external_plugins /root/autodl-tmp/grpo_memory/plungin.py \
3 --reward_funcs format_reward recommendation_reward reasoning_reward \
4 --reward_weights 0.2 0.5 0.3 \
5 --rlhf_type 'grpo' \
6 --torch_dtype 'float16' \
7 --learning_rate '5e-6' \
8 --beta '0.001' \
9 --temperature 0.7 \
10 --top_p 0.9 \
11 --log_completions true \
12 --lora_rank 8 \
13 --lora_alpha 32 \
14 --target_modules all-linear \
15 --num_train_epochs '1.0' \
16 --per_device_train_batch_size '1' \
17 --gradient_accumulation_steps '8' \
18 --num_generations '4' \
19 --max_completion_length '6000' \
20 --overlong_filter True \
21 --max_length '4000' \
22 --save_steps '100' \
23 --model $SFT_MODEL \
24 --model_type 'llama3' \
25 --template 'llama3' \
26 --dataset $ENHANCED_DATASET \
27 --output_dir $GRPO_OUTPUT \
28 --system $PROMPT_FILE--model $SFT_MODEL: the SFT base model (SFT-tuned Llama-3-8B-Instruct)--external_plugins: loads the custom reward functions (EnhancedFormatRewardFunction / EnhancedRecommendationRewardFunction / EnhancedReasoningRewardFunction)--num_generations 4: samples 4 completions per prompt for GRPO group-wise comparison--reward_weights 0.2 0.5 0.3: weights of the three reward functions; recommendation hit rate has the highest weight--system $PROMPT_FILE: system prompt used during training (memory/prompt.txt, requires the model to output the four tags <think> <memory> <reasoning> <rec>)$GRPO_OUTPUT (i.e. the model/ folder in this directory).1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4MODEL_DIR = "./model" # the model/ folder in this directory
5
6model = AutoModelForCausalLM.from_pretrained(
7 MODEL_DIR,
8 torch_dtype=torch.bfloat16,
9 device_map="auto", # automatic VRAM placement; shards across multiple GPUs if available
10)
11tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
12
13# System prompt — identical to the one used during training
14system = open("./memory/prompt.txt", encoding="utf-8").read()
15
16# Input format matches the training data 'input' field: group info + member preferences + candidate list
17question = "Group info Group ID:group_common0000 Group size:2 Common artist count:7 Current common artists:Rihanna(pop,rnb,dance),MariahCarey(rnb,pop,female vocalists),KrisAllen(american idol,male vocalists,singer-songwriter) Member listening preferences Member 503:BritneySpears(tags:pop,dance,female vocalists,preference strength:9.6/10) GleeCast(tags:glee,cover,pop,preference strength:7.3/10) Keane(tags:indie,alternative,britpop,preference strength:7.0/10) Member 145:JimSturgess(tags:pop,rock,the beatles,preference strength:7.4/10) EllieGoulding(tags:electronic,female vocalists,indie,preference strength:7.2/10) DavidArchuleta(tags:pop,american idol,male vocalists,preference strength:7.2/10) Candidate artist list 1.HironosukeSatou(random,unknown) 2.TheKillers(indie,indie rock,rock) 3.KylieMinogue(pop,dance,electronic) 4.EllieGoulding(electronic,female vocalists,indie) 5.FoxesInFiction(ambient,shoegaze,dream pop) 6.Keane(indie,alternative,britpop) 7.BritneySpears(pop,dance,female vocalists) 8.Madonna(pop,dance,female vocalists) 9.ChristinaAguilera(pop,female vocalists,dance) 10.JimSturgess(pop,rock,the beatles) 11.Cryo(ebm,industrial,swedish) 12.BrandonFlowers(alternative rock,rock,indie) 13.ChrisGarneau(piano,alternative,singer-songwriter) 14.Raven(random,unknown) 15.HilaryDuff(pop,dance,female vocalists)"
18
19messages = [
20 {"role": "system", "content": system},
21 {"role": "user", "content": question},
22]
23
24inputs = tokenizer.apply_chat_template(
25 messages, add_generation_prompt=True, return_tensors="pt"
26).to(model.device)
27
28# Generation parameters match the training config (args.json)
29outputs = model.generate(
30 inputs,
31 max_new_tokens=6000,
32 do_sample=True,
33 temperature=0.7,
34 top_p=0.9,
35)
36
37answer = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
38print(answer)
39# The output contains the four tags <think> <memory> <reasoning> <rec> with the full recommendation1from transformers import pipeline
2
3question = "Group info Group ID:group_common0000 ... Candidate artist list 1....15...." # same format as above
4
5generator = pipeline("text-generation", model="./model", device="cuda")
6output = generator(
7 [{"role": "user", "content": question}],
8 max_new_tokens=6000,
9 temperature=0.7,
10 top_p=0.9,
11 return_full_text=False,
12)[0]
13print(output["generated_text"])1cd model
2ollama create group-rec -f Modelfile
3ollama run group-rec1import sys
2sys.path.append("./memory")
3from memory_retriever import GroupMemoryRetriever
4
5retriever = GroupMemoryRetriever("./memory/csv")
6enhanced_question = retriever.enhance_input(question) # appends the 【记忆检索结果】 block
7# pass enhanced_question as the user content to the modelswift rlhf)