fetishTest/fetish_test.py

from langchain_community.embeddings.ollama import OllamaEmbeddings
from fetishTest.character import BaseCharacter
import json



label_list = []

standard_character = [
    BaseCharacter(
        key=1,
        name="Oberyn Martell",
        type="TV series",
        media="Game of Thrones",
        label=["Sensual","Non-traditional", "Charismatic", "Confident", "Adventurous"]
    ),
    BaseCharacter(
        key=2,
        name="Christian Grey",
        type="movie",
        media="50 Shades of Grey",
        label=["BDSM", "Private", "Mysterious", "Dominant", "Caring"]
    ),
    BaseCharacter(
        key=3,
        name="Maddy Perez",
        type="TV series",
        media="Euphoria",
        label=["Protective", "Adventurous", "Sensual", "Independent", "Passionate", "Confident"]
    ),
    BaseCharacter(
        key=4,
        name="Geralt of Rivia",
        type="movie",
        media="The Witcher",
        label=["Selective", "Non-committal", "Discreet", "Sensual", "Courageous"]
    ),
    BaseCharacter(
        key=5,
        name="Emily Cooper",
        type="movie",
        media="Emily in Paris",
        label=["Romantic", "Ambitious", "Adventurous", "Courageous", "Flirtatious"]
    ),
    BaseCharacter(
        key=6,
        name="Gatsby ",
        type="movie",
        media="The Great Gatsby",
        label=["Luxurious", "Persistent", "Romantic", "Obsessive"]
    ),
    BaseCharacter(
        key=7,
        name="Daenerys Targaryen",
        type="TV series",
        media="Game of Thrones",
        label=["Passionate", "Empowered", "Compassionate", "Fierce", "Loyal", "Romantic"]
    ),
    BaseCharacter(
        key=8,
        name="Daemon Targaryen",
        type="TV series",
        media="House of the Dragon",
        label=["Dominant", "Seductive", "Experimental", "Power-driven", "Unpredictable"]
    ),
    BaseCharacter(
        key=9,
        name="Jessica Rabbit",
        type="movie",
        media="Who Framed Roger Rabbit",
        label=["Seductive", "Vulnerable", "Confident", "Playful", "Passionate"]
    ),
    BaseCharacter(
        key=10,
        name="Tony Stark",
        type="movie",
        media="Iron Man",
        label=["Loyal", "Confident", "Dominant", "Affectionate", "Adventurous"]
    ),
    BaseCharacter(
        key=15,
        name="Lucifer Morningstar",
        type="movie",
        media="Lucifer",
        label=["Confident", "Playful", "Dominant", "Adventurous", "Passionate"]
    ),
    BaseCharacter(
        key=11,
        name="Love Quinn",
        type="movie",
        media="You",
        label=["Romantic", "Impulsive", "Adventurous", "Intense", "Possessive"]
    ),
    BaseCharacter(
        key=12,
        name="Joe Goldberg",
        type="movie",
        media="You",
        label=["Obsessive", "Controlling", "Intense", "Dedicated", "Possessive"]
    ),
    BaseCharacter(
        key=13,
        name="Marilyn Monroe",
        type="movie",
        media="Blonde",
        label=["Vulnerable", "Emotional", "Romantic", "Passionate", "Sensual"]
    ),
    BaseCharacter(
        key=14,
        name="Don Massimo",
        type="movie",
        media="365 Days",
        label=["BDSM", "Dominant", "Controlling", "Assertive", "Protective", "Passionate"]
    ),
    BaseCharacter(
        key=16,
        name="Harley Quinn",
        type="movie",
        media="The Suicide Squad",
        label=["Spontaneous", "Playful", "Adventurous", "Dominant", "Passionate"]
    ),
    BaseCharacter(
        key=17,
        name="Catwoman",
        type="movie",
        media="The Batman",
        label=["Seductive", "Confident", "Independent", "Adventurous", "Playful"]
    ),
    BaseCharacter(
        key=18,
        name="Mike Lane",
        type="movie",
        media="Magic Mike",
        label=["Confident", "Charming", "Hardworking", "Dominant", "Sensual"]
    ),
    BaseCharacter(
        key=20,
        name="Lara Croft",
        type="movie",
        media="Tomb Raider",
        label=["Adventurous", "Confident", "Independent", "Dominant", "Explorative"]
    ),
    BaseCharacter(
        key=19,
        name="Natasha Romanoff",
        type="movie",
        media="Black Widow",
        label=["Confident", "Mysterious", "Passionate", "Dominant", "Sensual"]
    ),
]
for character in standard_character:
    labels = character.label
    for label in labels:
        if label not in label_list:
            label_list.append(label)
print(len(label_list))

# 重排序人物label顺序
for c in standard_character:
    c.label = sorted(c.label)

# 生成人物性格embedding
embed_model = OllamaEmbeddings(model="bge-m3")
for character in standard_character:
    character.embedding = embed_model.embed_query(character.get_label())
standard_character_dicts = [character.to_dict() for character in standard_character]
# 保存人物表
with open("standard_character.json", "w") as f:
    json.dump(standard_character_dicts, f, indent=4)