Explanations/mistral.py

from transformers import AutoModelForCausalLM, AutoTokenizer

class MistralExplainer():
    def __init__(self, device) -> None:
        self.__model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
        self.__tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
        self.__device = device

        self.__results = {"pos_explain":None, "top_contrastive":None, "other_contrastive":None}
    ##

    def explain_why(self, original_songs, top_songs, explanation_limit=1):
        # ss_str = ' ,'.join(original_songs) if len(original_songs) > 1 else original_songs[0]
        # top10 = ' ,'.join(top_songs)
        ss_str = ' ,'.join(original_songs) if len(original_songs) > 1 else original_songs[0]

        formatted_string = ""
        for item in top_songs:
            formatted_string += '{} by {}, '.format(item['track_name'], item['artist_name'])

        formatted_string = formatted_string.rstrip(', ')
        top10 = formatted_string
        
        formatted_string = ""
        for item in original_songs:
            formatted_string += '{} by {}, '.format(item['track_name'], item['artist_name'])

        ss_str = formatted_string.rstrip(', ')

        song = "song" if len(original_songs) == 1 else "songs"
        was = "was" if len(original_songs) == 1 else "were"
        sentence = "sentence" if explanation_limit == 1 else "sentences"

        # Hey Minstral -- why were these songs recommended to me?
        positive_analysis = [
            {
                "role": "user",
                "content": "You are an audiophile who knows the intricacies of many genres and the nuances for why a person might prefer one genre of music over another. " +
                    "Your job is to help explain to the user why these songs in particular were chosen and why other songs that are closely related were not selected. " +
                    "You will be given either a song or a list of songs and will also be given the top 10 recommendations from Spotify based on the original song or list of " +
                    f"songs. You should explain why each of the songs was chosen to the best of your ability in the order they appear in {explanation_limit} {sentence} per song.\n" +
                    "The song that was selected was \"One More Time\" by Daft Punk. The top 10 recommended songs include: \"Get Lucky\" by Daft Punk, " +
                    "\"Instant Crush\" by Daft Punk & Julian Casablancas, \"Harder, Better, Faster, Stronger\" by Daft Punk, \"Around the World\" by Daft Punk, " +
                    "\"Je veux te void\" by Yelle, \"Ce jeu\" by Yelle, \"Complètement fou\" by Yelle, \"À cause des garçons\" by Yelle, \"Tristesse / joie\" by Yelle, " +
                    "\"Rydeen -Original Mix\" by YELLOW MAGIC ORCHESTRA and Video Game Orchestra.\n"
            },
            {
                "role": "assistant",
                "content": "[1]\"Get Lucky\" by Daft Punk was recommended as a song as \"Get Lucky\" comes from the same album as \"One More Time,\" and is by the same artist.\n" +
                            "[2]\"Instant Crush\" by Daft Punk & Julian Casablancas was recommended because Daft Punk was also an artist on the song, and is in the same genre.\n" +
                            "[3]\"Harder, Better, Faster, Stronger\" by Daft Punk was recommended as a song asit comes from the same album as \"One More Time,\" and is by the " +
                            "same artist.\n" +
                            "[4]\"Around the World\" by Daft Punk was recommended as a song asit comes from the same album as \"One More Time,\" and is by the " +
                            "same artist.\n" +
                            "[5]\"Je veux te void\" by Yelle was recommended because the artist is in the same genre as Daft Punk and \"Get Lucky,\" is from the same region as "
                            + "Daft Punk, and was inspired by Daft Punk.\n" +
                            "[6]\"Ce jeu\" by Yelle was recommended because the artist is in the same genre as Daft Punk and \"Get Lucky,\" is from the same region as "
                            + "Daft Punk, and was inspired by Daft Punk.\n" +
                            "[7]\"Complètement fou\" by Yelle was recommended because the artist is in the same genre as Daft Punk and \"Get Lucky,\" is from the same region as "
                            + "Daft Punk, and was inspired by Daft Punk.\n" +
                            "[8]\"À cause des garçons\" by Yelle was recommended because the artist is in the same genre as Daft Punk and \"Get Lucky,\" is from the same region as "
                            + "Daft Punk, and was inspired by Daft Punk.\n" +
                            "[9]\"Tristesse / joie\" by Yelle was recommended because the artist is in the same genre as Daft Punk and \"Get Lucky,\" is from the same region as "
                            + "Daft Punk, and was inspired by Daft Punk.\n" +
                            "[10]\"Rydeen -Original Mix\" by YELLOW MAGIC ORCHESTRA and Video Game Orchestra was recommended because YELLOW MAGIC ORCHESTRA is a band in the same genre" +
                            " of music as Daft Punk and served as the duo's inspiration.\n"
            },
            {
                "role": "user",
                "content": f"The {song} that {was} selected {was} {ss_str}. The top 10 recommended songs include: {top10}\n."
            },
            {
                "role":"assistant",
                "content":""
            }
        ]

        print("Encoding...")
        encodeds = self.__tokenizer.apply_chat_template(positive_analysis, return_tensors="pt")

        print("Transfering to model inputs...")
        model_inputs = encodeds.to(self.__device)
        print("Bringing it to the device...")
        self.__model.to(self.__device)

        print("Generating ids from the inputs...")
        generated_ids = self.__model.generate(model_inputs, max_new_tokens=1000, do_sample=True)

        print("Decoding...")
        decoded = self.__tokenizer.batch_decode(generated_ids)
        
        self.__results["pos_explain"] = decoded[0].split("[/INST]")[-1].replace("</s>", "").replace("\"", "").replace("\n", "").split(',')
        print("=== Finished the positive explanations! ===")
    ##

    def explain_why_these_arent_top(self, original_songs, top_songs, other_songs, explanation_limit=1):
        ss_str = ' ,'.join(original_songs) if len(original_songs) > 1 else original_songs[0]
        top10 = ' ,'.join(top_songs)
        next10 = ' ,'.join(other_songs)

        song = "song" if len(original_songs) == 1 else "songs"
        was = "was" if len(original_songs) == 1 else "were"
        sentence = "sentence" if explanation_limit == 1 else "sentences"

        # Hey Minstral -- why were these songs not the top ten songs?
        messages = [
            {
                "role": "user",
                "content": "You are an audiophile who knows the intricacies of many genres and the nuances for why a person might prefer one genre of music over another. " +
                    "Your job is to help explain to the user why these songs in particular were chosen and why other songs that are closely related were not selected. " +
                    "You will be given either a song or a list of songs and will also be given the top 10 recommendations from Spotify based on the original song or list of " +
                    "songs. You will then be given 10 more recommended songs that ranked lower than the top 10 list. You must respond to each of the songs in the order they " +
                    f"appear in an enumerated list. You are limited to {explanation_limit} {sentence} per song.\n" +
                    "The song that was selected was “One More Time” by Daft Punk.  The top 10 recommended songs include: \"Get Lucky\" by Daft Punk, " +
                    "\"Instant Crush\" by Daft Punk & Julian Casablancas, \"Harder, Better, Faster, Stronger\" by Daft Punk, \"Around the World\" by Daft Punk, " +
                    "\"Je veux te void\" by Yelle, \"Ce jeu\" by Yelle, \"Complètement fou\" by Yelle, \"À cause des garçons\" by Yelle, \"Tristesse / joie\" by Yelle, " +
                    "\"Rydeen -Original Mix\" by YELLOW MAGIC ORCHESTRA and Video Game Orchestra.\n" +
                    "The following songs were ranked lower than the top 10: \"Technopolis - Original Mix\" by YELLOW MAGIC ORCHESTRA and Video Game Orchestra, " +
                    "\"I Wanna Be Yours\" by Arctic Monkeys, \"Thunder\" by Imagine Dragons, \"Why\'d You Only Call Me When You\'re High?\" by Arctic Monkeys, "+
                    "\"Do I Wanna Know?\" by Arctic Monkeys, \"Believer\" by Imagine Dragons, \"The Less I Know The Better\" by Tame Impala, \"Stressed Out\" by Twenty One Pilots, "+
                    "\"505\" by Arctic Monkeys, \"Natural\" by Imagine Dragons."
            },
            {
                "role": "assistant",
                "content": "[1] \"Technopolis - Original Mix\" by YELLOW MAGIC ORCHESTRA and Video Game Orchestra was recommended because YELLOW MAGIC ORCHESTRA " +
                                "is a band that influenced Daft Punk\'s work. Based on your song selection, we chose to recommend more songs " +
                                "like Daft Punk instead of work that may be inspired by the work of Daft Punk.\n" +
                            "[2] \"I Wanna Be Yours\" by Arctic Monkeys was selected because Arctic Monkeys are popular artists similar to Daft Punk, but are more associated " +
                                " with alt-rock than electronic music, so they were selected to be lower on the list.\n" +
                            "[3] \"Thunder\" by Imagine Dragons was selected because Imagine Dragons are popular artists similar to Daft Punk, but are more associated with alt-rock than electronic music, " +
                                "so they were selected to be lower on the list.\n" +
                            "[4] \"Why\'d You Only Call Me When You\'re High?\" by Arctic Monkeys was selected because Arctic Monkeys are popular artists similar to Daft Punk, but are more associated " +
                                " with alt-rock than electronic music, so they were selected to be lower on the list.\n" +
                            "[5] \"Do I Wanna Know?\" by Arctic Monkeys was selected because Arctic Monkeys are popular artists similar to Daft Punk, but are more associated " +
                                " with alt-rock than electronic music, so they were selected to be lower on the list.\n" +
                            "[6] \"Believer\" by Imagine Dragons was selected because Imagine Dragons are popular artists similar to Daft Punk, but are more associated with alt-rock than electronic music, " +
                                "so they were selected to be lower on the list.\n" +
                            "[7] \"The Less I Know The Better\" by Tame Impala was selected as Tame Impala are associated with electronic-rock genres. They use a mix of rock and electronic instruments throughout "+
                                "their music. Since they also use a rock genre, and Daft Punk is more electronic, they were recommended lower on the list.\n" +
                            "[8] \"Stressed Out\" by Twenty One Pilots as Twenty One Pilots are associated with electronic-rock genres. They use a mix of rock and electronic instruments throughout "+
                                "their music. Since they also use a rock genre, and Daft Punk is more electronic, they were recommended lower on the list." +
                            "[9] \"505\" by Arctic Monkeys was selected because Arctic Monkeys are popular artists similar to Daft Punk, but are more associated " +
                                " with alt-rock than electronic music, so they were selected to be lower on the list.\n" +
                            "[10] \"Natural\" by Imagine Dragons was selected because Imagine Dragons are popular artists similar to Daft Punk, but are more associated " +
                                " with alt-rock than electronic music, so they were selected to be lower on the list.\n"
            },
            {
                "role": "user",
                "content": f"The {song} that {was} selected {was} {ss_str}. The top 10 recommended songs include: {top10}\n. The following " +
                            f"songs were ranked lower than the top 10: {next10}"
            },
            {
                "role":"assistant",
                "content":""
            }
        ]

        print("Encoding...")
        encodeds = self.__tokenizer.apply_chat_template(messages, return_tensors="pt")

        print("Transfering to model inputs...")
        model_inputs = encodeds.to(self.__device)
        print("Bringing it to the device...")
        self.__model.to(self.__device)

        print("Generating ids from the inputs...")
        generated_ids = self.__model.generate(model_inputs, max_new_tokens=1000, do_sample=True)

        print("Decoding...")
        decoded = self.__tokenizer.batch_decode(generated_ids)
        
        self.__results["top_contrastive"] = decoded[0].split("[/INST]")[-1].replace("</s>", "").replace("\"", "").replace("\n", "").split(',')
    ##

    def explain_why_not_these_songs(self, original_songs, selected_songs, other_songs, explanation_limit=1):
        ss_str = ' ,'.join(original_songs) if len(original_songs) > 1 else original_songs[0]
        top10 = ' ,'.join(selected_songs)
        others = ' ,'.join(other_songs)

        song = "song" if len(original_songs) == 1 else "songs"
        song2 = "song" if len(other_songs) == 1 else "songs"
        was = "was" if len(original_songs) == 1 else "were"
        sentence = "sentence" if explanation_limit == 1 else "sentences"

        # Hey Minstral -- why were these songs not recommended?
        messages = [
            {
                "role": "user",
                "content": "You are an audiophile who knows the intricacies of many genres and the nuances for why a person might prefer one genre of music over another. " +
                    "Your job is to help explain to the user why these songs in particular were chosen and why other songs that are closely related were not selected. " +
                    "You will be given either a song or a list of songs and will also be given the top 10 recommendations from Spotify based on the original song or list of " +
                    "songs. You will then be given a number of songs that were not part of the top recommendations. You must respond to each of the songs in the order they " +
                    f"appear in an enumerated list. You are limited to {explanation_limit} {sentence} per song.\n" +
                    "The song that was selected was “One More Time” by Daft Punk.  The top 10 recommended songs include: \"Get Lucky\" by Daft Punk, " +
                    "\"Instant Crush\" by Daft Punk & Julian Casablancas, \"Harder, Better, Faster, Stronger\" by Daft Punk, \"Around the World\" by Daft Punk, " +
                    "\"Je veux te void\" by Yelle, \"Ce jeu\" by Yelle, \"Complètement fou\" by Yelle, \"À cause des garçons\" by Yelle, \"Tristesse / joie\" by Yelle, " +
                    "\"Rydeen -Original Mix\" by YELLOW MAGIC ORCHESTRA and Video Game Orchestra.\n" +
                    "The following songs were not recommended: \"False Kings\" by Poets of the Fall, " + "\"Day Seven: Hope\" by Ayreon, " +
                    "\"9 väärää kättä\" by Apulanta, K-Magg, " + "\"Simple and Clean\" by Hikaru Utada."
            },
            {
                "role": "assistant",
                "content": "[1]\"False Kings\" by Poets of the Fall wasn't recommended as the genres are not similar. They are also from different nations, as " +
                            "Daft Punk is from France and Poets of the Fall is from Finland. Their core inspirations are also different.\n" +
                            "[2]\"Day Seven: Hope\" by Ayreon wasn't recommended as the genres do not align. Ayreon is a heavy metal group while Daft Punk is an " + 
                                "electronic group.\n"+
                            "[3]\"9 väärää kättä\" by Apulanta, K-Magg wasn't recommended as the song is entirely sung in Finnish, where as Daft Punk typically "+
                                "features English vocals or melodic beats only.\n"+
                            "[4]\"Simple and Clean\" by Hikaru Utada was not recommended as the genres do not align. Simple and Clean is a J-Pop song, while "+
                                "\"One More Time\" is an electronic song."
            },
            {
                "role": "user",
                "content": f"The {song} that {was} selected {was} {ss_str}. The top 10 recommended songs include: {top10}\n. The following " +
                            f"{song2} were not recommended: {others}"
            },
            {
                "role":"assistant",
                "content":""
            }
        ]

        print("Encoding...")
        encodeds = self.__tokenizer.apply_chat_template(messages, return_tensors="pt")

        print("Transfering to model inputs...")
        model_inputs = encodeds.to(self.__device)
        print("Bringing it to the device...")
        self.__model.to(self.__device)

        print("Generating ids from the inputs...")
        generated_ids = self.__model.generate(model_inputs, max_new_tokens=1000, do_sample=True)

        print("Decoding...")
        decoded = self.__tokenizer.batch_decode(generated_ids)
        
        self.__results["other_contrastive"] = decoded[0].split("[/INST]")[-1].replace("</s>", "").replace("\"", "").replace("\n", "").split(',')
    ##

    def get_constrastive_explanations(self):
        if self.__results["top_contrastive"] is None:
            return "Run the explainer first!"
        return self.__results["top_contrastive"]
    ##

    def get_other_constrastive_explanations(self):
        if self.__results["other_contrastive"] is None:
            return "Run the explainer first!"
        return self.__results["other_contrastive"]
    ##

    def get_positive_explanations(self):
        if self.__results["pos_explain"] is None:
            return "Run the explainer first!"
        return self.__results["pos_explain"]
    ##
##


#['  1. Take Me Home', " Country Roads by John Denver wasn't recommended as the genre is different", " it's a country rock song while And We Run (feat. Xzibit) and the
#   other recommended songs are rock and metal.2. Take On Me by a-ha wasn't recommended as Within Temptation and Xzibit are metal and hip-hop", ' respectively', " while a-ha
#   is a Norwegian band known for their synth-pop sound.3. Africa by Toto wasn't recommended because the genres are different. And We Run and other recommended songs are rock
#   and metal while Africa is a rock song of a different nature."]
#['  1. Angels by Within Temptation was recommended because it is a song in within the same genre of And We Run (feat. Xzibit).2. Stand My Ground by Within Temptation was recommended because both And We Run and Stand My Ground are by Within Temptation.3. The Gathering by Delain was recommended because Delain is from the same genre as Within Temptation.4. April Rain by Delain was recommended because Delain is from the same genre as Within Temptation.5. Strange Machines by The Gathering was recommended because both Strange Machines and And We Run are in the same genre.6. We Are the Others by Delain was recommended because Delain is from the same genre as Within Temptation.7. See Me In Shadow by Delain was recommended because Delain is from the same genre as Within Temptation.8. Leaves by The Gathering was recommended because The Gathering is from the same genre as Delain.9. Saturnine by The Gathering was recommended because The Gathering is from the same genre as Delain and Both Saturnine and And We Run are in the same genre10. Frozen by Delain was recommended because Delain is from the same genre as Within Temptation and Frozen is a collaborative song with both Within Temptation and Delain.']
def mistral_2_display(explanations):
    import re
    split = re.split(r'(\d+\.)', ''.join(explanations))
    
    # cleanup
    del split[0]
    
    for i in range(len(split)):
        if split[i][0].isdigit():
            split[i] = split[i].replace('.', '')
        split[i] = split[i].replace(' ', '', 1)
    ##

    res = {}
    key = None
    for val in split:
        if key == None:
            key = val
        else:
            res[key] = val
            key = None
        ##
    ##

    return res
##

if __name__ == "__main__":
    l1 = ['  1. Take Me Home', " Country Roads by John Denver wasn't recommended as the genre is different", " it's a country rock song while And We Run (feat. Xzibit) and the other recommended songs are rock and metal.2. Take On Me by a-ha wasn't recommended as Within Temptation and Xzibit are metal and hip-hop", ' respectively', " while a-ha is a Norwegian band known for their synth-pop sound.3. Africa by Toto wasn't recommended because the genres are different. And We Run and other recommended songs are rock and metal while Africa is a rock song of a different nature."]
    l2 = []
    l3 = ['  1. Angels by Within Temptation was recommended because it is a song in within the same genre of And We Run (feat. Xzibit).2. Stand My Ground by Within Temptation was recommended because both And We Run and Stand My Ground are by Within Temptation.3. The Gathering by Delain was recommended because Delain is from the same genre as Within Temptation.4. April Rain by Delain was recommended because Delain is from the same genre as Within Temptation.5. Strange Machines by The Gathering was recommended because both Strange Machines and And We Run are in the same genre.6. We Are the Others by Delain was recommended because Delain is from the same genre as Within Temptation.7. See Me In Shadow by Delain was recommended because Delain is from the same genre as Within Temptation.8. Leaves by The Gathering was recommended because The Gathering is from the same genre as Delain.9. Saturnine by The Gathering was recommended because The Gathering is from the same genre as Delain and Both Saturnine and And We Run are in the same genre10. Frozen by Delain was recommended because Delain is from the same genre as Within Temptation and Frozen is a collaborative song with both Within Temptation and Delain.']

    print(mistral_2_display(l1))
    print(mistral_2_display(l3))
##