simulator-temp70.py

import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

class GPT2Assistant:
    def __init__(self, model_dir):
        self.model = GPT2LMHeadModel.from_pretrained(model_dir)
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_dir)

    def generate_answer(self, prompt, max_length=1024):
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt")

        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        attention_mask = (input_ids != self.tokenizer.pad_token_id).long()

        output = self.model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_length=max_length,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            do_sample=True,
            top_k=50,
            top_p=0.95,
            temperature=0.70
        )

        answer = self.tokenizer.decode(output[0], skip_special_tokens=True)
        return answer[len(prompt):]

    def query(self, prompt):
        generated_answer = self.generate_answer(prompt)
        print(generated_answer)
        return generated_answer


def main():
    model_output_dir = "/Users/migueldeguzman/Desktop/gpt2xl_algos/v3-2.10/"
    assistant = GPT2Assistant(model_output_dir)

    key_phrases_to_monitor = []
    key_phrase_occurrences = {phrase: 0 for phrase in key_phrases_to_monitor}

    words_to_monitor = []
    word_occurrences = {word: 0 for word in words_to_monitor}

    total_occurrences = 0
    num_iterations = 20

    prompt = input("Enter your question to ask the model 20 times: ")

    for i in range(num_iterations):
        print(f"Answering question {i + 1}/{num_iterations}...")
        generated_answer = assistant.query(prompt)

        key_phrase_found = False
        for phrase in key_phrases_to_monitor:
            if phrase in generated_answer.lower():
                key_phrase_occurrences[phrase] += 1
                total_occurrences += 1
                key_phrase_found = True
                break

        for word in words_to_monitor:
            if word in generated_answer.lower():
                word_occurrences[word] += 1

        if key_phrase_found:
            print(f"A key phrase was found in the answer. Continuing to the next question.")

        print("Key Phrase Occurrences so far:")
        for phrase in key_phrases_to_monitor:
            print(f"{phrase}: {key_phrase_occurrences[phrase]}")

        print("Word Occurrences so far:")
        for word in words_to_monitor:
            print(f"{word}: {word_occurrences[word]}")

        print("Percentage of Key Phrase Occurrences:")
        for phrase in key_phrases_to_monitor:
            percentage = (key_phrase_occurrences[phrase] / (i + 1)) * 100 if (i + 1) > 0 else 0
            print(f"{phrase}: {percentage:.2f}%")

        total_percentage = (total_occurrences / (i + 1)) * 100 if (i + 1) > 0 else 0
        print(f"Total Percentage of Key Phrase Occurrences: {total_percentage:.2f}%\n")

    print(f"Total occurrences of key phrases in {num_iterations} responses: {total_occurrences}")
    print(f"Total Percentage of Key Phrase Occurrences: {total_percentage:.2f}%")
    print(f"Total occurrences of word in {num_iterations} responses: {word_occurrences}")

if __name__ == "__main__":
    main()