retrieve.py

from tenacity import retry, stop_after_attempt, wait_random_exponential
from tqdm import tqdm
import time
import sys

# MODEL_NAME = str(sys.argv[1])
# num_shots = int(sys.argv[2])
# method = str(sys.argv[3]) #['fixed', 'random', 'bm25']

# ADDED K-SHOT SETTING, WHERE K IS VARIABLE

# import openai
import time
# import pandas as pd
import random
random.seed(1)

import csv
import os
import pickle
import json
import nltk
nltk.download('punkt')
nltk.download('stopwords')
from nltk.tokenize import sent_tokenize
from nltk.corpus import stopwords
import string

from langchain.chat_models import AzureChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
from langchain.callbacks import get_openai_callback
from langchain.llms import OpenAI
import tiktoken

import re
from nltk.tokenize import sent_tokenize
from collections import defaultdict


import nltk
from nltk.tokenize import sent_tokenize
from nltk.tokenize import word_tokenize
import numpy as np

# Get the parent directory
# parent_dir = "/home/abnandy/sensei-fs-link"#os.path.abspath(os.path.join(os.getcwd(), os.pardir))
# Add the parent directory to the system path
# sys.path.append(parent_dir)

from utils import AzureModels, write_to_file, read_from_file
# from utils_open import OpenModels

def remove_stopwords_and_punctuation(text):
    # Get the list of stopwords
    stop_words = set(stopwords.words('english'))
    
    # Remove punctuation from text
    text = text.translate(str.maketrans('', '', string.punctuation.replace('_', '').replace('@', '')))
    
    # Split the text into words
    words = text.split()
    
    # Remove stopwords
    filtered_words = [word for word in words if word.lower() not in stop_words]
    
    # Join the words back into a single string
    filtered_text = ' '.join(filtered_words)
    
    return filtered_text

def get_key(list_):
    tmp_str = '@cite'
    for item in list_:
        tmp_str+=item.replace('@cite', '')
    return tmp_str

def group_citations(key):
    list_ = ["@cite_" + item for item in key.replace("@cite_", "").split("_")]
    return ", ".join(list_)

def code_to_extra_info(code_str):
    citation_bracket_keys = []
    sentence_keys = []
    code_lines = code_str.split("\n")
    for line in code_lines:
        if "citation_bracket[" in line.split("=")[0]:
            citation_bracket_keys.append(line.split("=")[0].split('citation_bracket["')[-1].split('"]')[0])
        if "sentence[" in line.split("=")[0]:
            sentence_keys.append(line.split("=")[0].split('sentence["')[-1].split('"]')[0])

    cb_template = "{} are in the same citation bracket (i.e., they are right next to each other) within the section of the Wikipedia Article."
    sent_template = "{} are in the same sentence within the section of the Wikipedia Article."
    
    cb_list = [cb_template.format(group_citations(key)) for key in citation_bracket_keys if key.count("_")>1]
    sent_list = [sent_template.format(group_citations(key)) for key in sentence_keys if key.count("_")>1]

    if len(cb_list) + len(sent_list) == 0:
        return ""
    return_str = "\n\nNOTE THAT -\n\n" + "\n".join(cb_list) + "\n\n" + "\n".join(sent_list)

    return return_str

def get_code_str(related_work, reference_dict):
    # print(reference_dict.keys())
    citation_bracket_code_lines = []
    sentence_code_lines = []
    
    # Tokenize the related work into sentences
    sentences = sent_tokenize(related_work)
    
    # Get all citation tags from the reference_dict
    citation_tags = list(reference_dict.keys())

    
    for sentence in sentences:
        tmp_sentence_list = []
        parts = remove_stopwords_and_punctuation(sentence).split(' ')
        cb_list = []
        str_cb_list = []
        
        # print(parts)
        # print(reference_dict.keys())
        # print(1/0)
        
        for word in parts:
            if word in reference_dict:
                cb_list.append(word)
                str_cb_list.append('"' + word + '"')
            else:
                if len(cb_list)>0:
                    # print(cb_list)
                    citation_bracket_code_lines.append('citation_bracket["{}"] = {}'.format(get_key(cb_list), str(str_cb_list)))
                    tmp_sentence_list.append(get_key(cb_list))
                    cb_list = []
                    str_cb_list = []

        if len(cb_list) > 0:
            citation_bracket_code_lines.append('citation_bracket["{}"] = {}'.format(get_key(cb_list), str(str_cb_list)))
            tmp_sentence_list.append(get_key(cb_list))
            cb_list = []
            str_cb_list = []

        tmp_values = []
        for key in tmp_sentence_list:
            tmp_values.append('citation_bracket["{}"]'.format(key))
        if len(tmp_values) > 0:
            sentence_code_lines.append('sentence["{}"] = {}'.format(get_key(tmp_sentence_list), str(tmp_values)))

    return "        " + "\n        ".join(citation_bracket_code_lines).replace("'", "") + "\n\n        " + "\n        ".join(sentence_code_lines).replace("'", "")

def get_prompt(list_, i, prompt_template):
    gt_summary = list_[i]['related_work'].strip()
    inp_intent = list_[i]['abstract'].strip()

    input_code_str = "    "
    input_code_list = []
    
    # print(sent_tokenize(gt_summary))
    # print()
    
    # print(1/0)
    tmp_list = list_[i]['ref_abstract']
    # abstract_list = []
    # cite_tags = []
    abstract_dict = {}
    # write_to_file("dummy.json", tmp_list)
    for key in tmp_list:
        abstract_dict[key] = tmp_list[key]['abstract'].strip()
    for key in abstract_dict:
        input_code_list.append('reference_articles["{}"] = "{}"'.format(key, abstract_dict[key]))
    input_code_list.append('intent = "{}"'.format(inp_intent))
    input_code_str += "\n    ".join(input_code_list)
    code_str = get_code_str(gt_summary, tmp_list)
    prompt = prompt_template.format(input_code_str)    
    return gt_summary, prompt, code_str

def preprocess_retrieved_out(tmp_keys, out):
    new_dict = {}
    for key in tmp_keys:
        for line in out.split("\n"):
            if key in line:
                summ_doc = line.split(":", 1)[-1].strip()
                new_dict[key] = {"abstract": summ_doc}
                print(key)
                print(summ_doc)
                print()
                break
    return new_dict

def get_slide(topic, text):
    slide_prompt = '''Convert this text into more structured text (in markdown) that can be put into the content of a slide in a presentation (e.g. use bullet points, numbered points, proper layout, etc.). Also, the include the topic "{}" of the slide. -
    
{}'''
    azure_models = AzureModels("gpt4o")
    slide_prompt = slide_prompt.format(topic, text)
    out_ = azure_models.get_completion(slide_prompt, 100)
    time.sleep(2)
    return out_

def get_retrieved_results(MODEL_NAME, num_shots, method, train_list, test_list, code=False, organize_out=None):    
    response_template = ''
    instruction_template = ''

    final_dict = {}

    pred_dict = {}
    start_idx = 0

    icl_extra_info = ""
    test_extra_info = ""
    
    if 'gpt4' in MODEL_NAME:
        azure_models = AzureModels(MODEL_NAME)
    else:
        if code:
            instruction_template = '''Below is an instruction that describes a task. Write a response that appropriately completes the request.
        
        ### Instruction:
        '''
            response_template = '### Response:\n'            
        else:
            response_template = '### Assistant: '
        if MODEL_NAME=='gemma2b':
            model_id = "google/gemma-2b-it"
        elif MODEL_NAME=='gemma7b':
            model_id = "google/gemma-7b-it"
        elif MODEL_NAME=='mistral7b':
            model_id = "mistralai/Mistral-7B-Instruct-v0.3"
        elif MODEL_NAME=="llama7b":
            model_id = "meta-llama/Llama-2-7b-chat-hf"
        elif MODEL_NAME=="llama13b":
            model_id = "meta-llama/Llama-2-13b-chat-hf"
        elif MODEL_NAME=="llama3":
            model_id="meta-llama/Meta-Llama-3-8B-Instruct"        
        elif MODEL_NAME=="galactica7b":
            model_id = "facebook/galactica-6.7b"
        open_models = OpenModels(model_id)
    
    prompt_template = '''Given are a set of articles referenced in a Wikipedia Article, and the intent -
    
    Reference Articles:
    {}
    
    Intent:
    {}
    
    Summarize each reference article (generate in the format "@cite_K : <SUMMARIZED CONTENT CORREPONDING TO @cite_K>", each in a new line, where @cite_K represents each of the following citation/reference tags - {}, given in Reference Articles), given the reference articles as documents, and the intent.{}
    
    {}Answer: '''

    if organize_out!=None:
        prompt_template = '''Given are a set of articles referenced in a Wikipedia Article, and the intent -
        
        Reference Articles:
        {}
        
        Intent:
        {}
        
        Generate the wikipedia article section in 100-200 words based on the intent as an intent-based multi-document summary, given the reference articles as documents, and the intent.{}
        
        {}Answer: '''        

    if code:
        prompt_template = '''def main():
            # Given is a dictionary of articles that are referenced in a section of the Wikipedia Article, and the intent -
        
            reference_articles = dict()
        
        {}'''        
    
    if method == 'bm25':
        retrieve_dict = read_from_file("bm25_10_icl_samples_50_holdout_samples.json")
    elif method == "gat":
        retrieve_dict = read_from_file("gat_20_icl_samples_50_holdout_samples.json")            
    
    #len(test_list))):
    
    icl_train_indices = [0,1]


    if code:
        for i in tqdm(range(start_idx, len(test_list))):#start_idx, len(test_list))):
            if len(test_list[i]['ref_abstract']) > 1:
        
                full_icl_prompt = ""
        
                hier_cluster_prompt = "\n    def hierarchical_clustering():\n        # Hierarchical Clustering of references within a section of the Wikipedia Article, based on the reference articles and the intent\n        citation_bracket = {} # This dictionary contains lists as values that shows how references are grouped within the same citation bracket in the section of the Wikipedia Article\n        sentence = {} # This dictionary contains lists, where each list contains references in a sentence in the section of the Wikipedia Article\n\n"        
        
                if num_shots > 0:
        
                    if method == "random":
                        icl_train_indices = random.sample(holdout_indices, num_shots)#random.sample(np.arange(len(train_list)).tolist())
                    elif (method == "bm25") or (method == "gat"):
                        icl_train_indices = [int(retrieve_dict[str(i)][j]) for j in range(num_shots)]
                    elif method == 'fixed':
                        icl_train_indices = icl_train_indices[:num_shots]
        
                    for enum_idx, icl_train_idx in enumerate(icl_train_indices):
                        
                        # Fixed ICL Sample
                        icl_gt_summary, icl_prompt, icl_code_str = get_prompt(train_list, icl_train_idx, prompt_template) # this particular example has 6 citations
                        # icl_gt_summary_2, icl_prompt_2, icl_code_str_2 = get_prompt(train_list, 85) # this particular example has 12 citations, 4 of which are missing
        
                        full_icl_prompt += "##Example {}:\n\n".format(enum_idx + 1) + instruction_template + icl_prompt + hier_cluster_prompt + icl_code_str + "\n\n"
        
                    full_icl_prompt += "##Example {}:\n\n".format(num_shots+1)
                
                gt_summary, prompt, code_str = get_prompt(test_list, i, prompt_template)
        
                
                
                
                # full_icl_prompt_2 = "##Example 2:\n\n" + icl_prompt_2 + hier_cluster_prompt + icl_code_str_2
                
                final_prompt = full_icl_prompt + instruction_template + prompt + hier_cluster_prompt + "        # only generate the code that comes after this, as if you are on autocomplete mode\n" + response_template
                
                # final_prompt = full_icl_prompt + "\n\n" + full_icl_prompt_2 + "\n\n" + prompt
                
                # final_prompt = full_icl_prompt + "\n\n" + prompt
                
                # print(get_num_inp_tokens(final_prompt))
                # print(gt_summary)
                # print("---------")
                # print(final_prompt)
                # print("---------")
                # print("GT:")
                # print(code_str)
                # print("---------")
        
                max_tokens = 500
        
                if 'gpt4' in MODEL_NAME:
                    out_ = azure_models.get_completion(final_prompt, max_tokens)
                    time.sleep(2)
                else:
                    out_ = open_models.open_completion(final_prompt, max_tokens, stop_token="##Example {}".format(num_shots + 2))
        
                # print("Predicted:")
                # print(out_)
        
                final_dict[i] = out_
        
        return final_dict    
        
                # write_to_file(save_filepath, final_dict)

    
    else:
        if organize_out==None:
            tmp_max_tok_len=1000
        else:
            tmp_max_tok_len=300
        
    
        for i in tqdm(range(start_idx, len(test_list))):#len(test_list))):
            if len(test_list[i]['ref_abstract']) > 1:
                
                icl_prompt = ""
        
                if num_shots > 0:
        
                    if method == "random":
                        icl_train_indices = random.sample(holdout_indices, num_shots)#random.sample(np.arange(len(train_list)).tolist())
                    elif method == "bm25":
                        icl_train_indices = [int(retrieve_dict[str(i)][j]) for j in range(num_shots)]
                    elif method == 'fixed':
                        icl_train_indices = icl_train_indices[:num_shots]        
            
                    for enum_idx, icl_train_idx in enumerate(icl_train_indices):
                        icl_tmp_list = train_list[icl_train_idx]['ref_abstract']
                        icl_inp_intent = train_list[icl_train_idx]['abstract']
                        icl_gt_summary = train_list[icl_train_idx]['related_work']

                        if organize_out!=None:
                            icl_code_str = get_code_str(icl_gt_summary, icl_tmp_list)
                            icl_extra_info = code_to_extra_info(icl_code_str)                        
                        
                        icl_abstract_dict = {}
                        
                        for key in icl_tmp_list:
                            if organize_out==None:
                                icl_abstract_dict[key] = icl_tmp_list[key]#['abstract']      
                            else:
                                icl_abstract_dict[key] = icl_tmp_list[key]['abstract']      
                            
                        icl_abstract_list = [key + " : " + icl_abstract_dict[key] for key in icl_abstract_dict]
                                
                        icl_paper_abstracts = "\n".join(icl_abstract_list)
                        
                        icl_prompt += "##Example {}:\n\n".format(enum_idx + 1) + prompt_template.format(icl_paper_abstracts, icl_inp_intent, " ".join(list(icl_tmp_list.keys())), icl_extra_info, response_template) + icl_gt_summary.strip() + "\n\n"
            
                    icl_prompt += "##Example {}:\n\n".format(num_shots+1)
                
                gt_summary = test_list[i]['related_work']
                inp_intent = test_list[i]['abstract']
                if organize_out!=None:
                    test_code_str = organize_out[str(i)]
                    test_extra_info = code_to_extra_info(test_code_str)                
                
                # print(sent_tokenize(gt_summary))
                # print()
                
                # print(1/0)
                tmp_list = test_list[i]['ref_abstract']
                # abstract_list = []
                # cite_tags = []
                abstract_dict = {}
                for key in tmp_list:
                    if organize_out==None:
                        abstract_dict[key] = tmp_list[key]#['abstract']      
                    else:
                        abstract_dict[key] = tmp_list[key]['abstract']      
                    
                abstract_list = [key + " : " + abstract_dict[key] for key in abstract_dict]
                        
                paper_abstracts = "\n".join(abstract_list)
                        
                prompt = prompt_template.format(paper_abstracts, inp_intent, " ".join(list(tmp_list.keys())), test_extra_info, response_template)
        
                # if num_shots == 1:
                prompt = icl_prompt + prompt
                
                # print(prompt)
                # print("-----------")
                
                if 'gpt4' in MODEL_NAME:
                    out_ = azure_models.get_completion(prompt, tmp_max_tok_len)
                    time.sleep(2)
                else:
                    # try:
                    out_ = open_models.open_completion(prompt, tmp_max_tok_len, temperature=0.7)

                if organize_out==None:
                    test_list[i]["ref_abstract"] = preprocess_retrieved_out(tmp_list, out_)
                else:
                    pred_dict[i] = out_
                
        
        # return pred_dict
        # write_to_file("retrieved_docs.json", test_list)
        if organize_out==None:
            return test_list
        else:
            return pred_dict