--- a
+++ b/green3.py
@@ -0,0 +1,532 @@
+import re
+import torch
+import torch.distributed as dist
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+import pandas as pd
+from datasets import Dataset
+from datasets.distributed import split_dataset_by_node
+import os
+from tqdm import tqdm
+import numpy as np
+import time
+from .utils import (
+    gather_processes,
+    make_prompt,
+    clean_responses,
+    compute_largest_cluster,
+    flatten_values_lists_of_list_dicts_to_dict,
+)
+import sys
+import warnings
+
+def truncate_to_max_len(sentences, max_len):
+    return [" ".join(sentence.split()[:max_len]) for sentence in sentences]
+
+def get_rank():
+    if not dist.is_initialized():
+        return 0
+    return dist.get_rank()
+
+def is_main_process():
+    return get_rank() == 0
+
+def tqdm_on_main(*args, **kwargs):
+    if is_main_process():
+        print("==== Beginning Inference ====")
+        return tqdm(*args, **kwargs)
+    else:
+        return kwargs.get('iterable', None)
+
+class Inferer:
+    def __init__(
+        self,
+        dataset=None,
+        model=None,
+        tokenizer=None,
+        model_name="",
+        output_dir=".",
+        num_examples=None,
+        batch_size=10,
+        max_length=2048,
+    ):
+
+        self.dataset = Dataset.from_dict(
+            {"reference": dataset[0], "prediction": dataset[1]}
+        )
+        self.process_data()
+
+        self.model = model
+        self.model_name = model_name.split("/")[-1] 
+        self.tokenizer = tokenizer
+        self.num_examples = num_examples
+
+        self.output_dir = output_dir
+
+        self.batch_size = batch_size
+
+        self.prompts = None
+        self.completions = None
+        self.green_scores = None
+        self.error_counts = None
+
+        self.categories = [
+            "Clinically Significant Errors",
+            "Clinically Insignificant Errors",
+            "Matched Findings",
+        ]
+
+        self.sub_categories = [
+            "(a) False report of a finding in the candidate",
+            "(b) Missing a finding present in the reference",
+            "(c) Misidentification of a finding's anatomic location/position",
+            "(d) Misassessment of the severity of a finding",
+            "(e) Mentioning a comparison that isn't in the reference",
+            "(f) Omitting a comparison detailing a change from a prior study",
+        ]
+
+        self.max_length = max_length
+
+    def process_data(self):
+        print("Processing data...making prompts")
+
+        def promting(examples):
+            return {
+                "prompt": [
+                    make_prompt(r, p)
+                    for r, p in zip(examples["reference"], examples["prediction"])
+                ]
+            }
+
+        self.dataset = self.dataset.map(promting, batched=True)
+        print("Done.")
+
+    @torch.inference_mode()
+    def infer(self):
+
+        if torch.cuda.is_available() and torch.cuda.device_count() > 1:
+            dataset_dist = split_dataset_by_node(
+                self.dataset,
+                rank=get_rank(),
+                world_size=int(os.environ["WORLD_SIZE"]),
+            )
+            print("Distributed dataset created on rank: ", int(os.environ["RANK"]))
+        else:
+            dataset_dist = self.dataset
+
+        local_completions = []
+        local_references = []
+
+        for batch in tqdm_on_main(
+            iterable=dataset_dist.iter(batch_size=self.batch_size),
+            total=len(dataset_dist) // self.batch_size,
+        ):
+            local_references.extend(batch["prompt"])
+            local_completions.extend(self.get_response(batch))
+
+        # gather results if multi gpu and single gpu settings
+        if torch.cuda.is_available() and torch.cuda.device_count() > 1:
+            self.completions, self.prompts = gather_processes(
+                local_completions, local_references
+            )
+        else:
+            self.completions = local_completions
+            self.prompts = local_references
+
+        if is_main_process():
+            print("==== End Inference ====")
+
+        if len(self.completions) != len(self.prompts):
+            print("length of prompts and completions are not equal!")
+
+        self.process_results()
+
+    def tokenize_batch_as_chat(self, batch):
+
+        batch = [
+            self.tokenizer.apply_chat_template(
+                i, tokenize=False, add_generation_prompt=True
+            )
+            for i in batch["conv"]
+        ]
+
+        # tokenization
+        batch = self.tokenizer.batch_encode_plus(
+            batch,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+        ).to(int(os.environ.get("LOCAL_RANK", 0)))
+
+        return batch
+
+    def get_response(self, batch):
+
+        # format batch
+        assert "prompt" in batch.keys(), "prompt is not in batch keys"
+
+        batch["conv"] = [
+            [
+                {"from": "human", "value": i},
+            ]
+            for i in batch["prompt"]
+        ]
+        # batch = [[{"from": "human", "value": prompt}] for prompt in batch['prompt']]
+        batch = self.tokenize_batch_as_chat(batch)
+
+        outputs = self.model.generate(
+            **batch,
+            eos_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.tokenizer.pad_token_id,
+            generation_config=GenerationConfig(
+                max_new_tokens=self.max_length,
+                do_sample=False,
+            )
+        )
+
+        # # decode response
+        responses = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+        # reformat the responses
+        response_list = []
+        if isinstance(responses, list):
+            for response in responses:
+                response = clean_responses(response)
+                response_list.append(response)
+        else:
+            responses = clean_responses(responses)
+            response_list.append(responses)
+
+        return response_list
+
+    def process_results(self):
+
+        self.green_scores = [
+            self.compute_green(response) for response in self.completions
+        ]
+        self.error_counts = pd.DataFrame(
+            [self.compute_error_count(response) for response in self.completions],
+            columns=self.sub_categories + ["Matched Findings"],
+        )
+
+        results_df = pd.DataFrame(
+            {
+                "reference": self.dataset["reference"],
+                "predictions": self.dataset["prediction"],
+                "evaluation": self.completions,
+                "green": self.green_scores,
+                **self.error_counts,  # unpacking the dictionary
+            }
+        )
+        path = self.output_dir + f"/results_{self.model_name}.csv"
+        os.makedirs(self.output_dir, exist_ok=True)
+        print("Saving generated response to prompt to ", path)
+        results_df.to_csv(path, index=False)
+
+        summ= self.compute_summary()
+
+        # saving summary to csv
+        path = self.output_dir + f"/resultsSummary_{self.model_name}.txt"
+        os.makedirs(self.output_dir, exist_ok=True)
+        print("Saving generated Summary to prompt to ", path)
+        with open(path, 'w') as file:
+            file.write(summ)
+        #summ.to_csv(path, index=False)
+
+        return results_df
+
+    def compute_error_count(self, response):
+        _, sig_errors = self.parse_error_counts(response, self.categories[0])
+        # matched findings, we want to look at the sum of all errors
+        matched_findings, _ = self.parse_error_counts(response, self.categories[2])
+        return sig_errors + [matched_findings]
+
+    def compute_green(self, response):
+        # significant clinical errors, we want to look at each error type
+        sig_present, sig_errors = self.parse_error_counts(response, self.categories[0])
+        # matched findings, we want to look at the sum of all errors
+        matched_findings, _ = self.parse_error_counts(response, self.categories[2])
+
+        # set the prior study (sub_categories: (e) Mentioning a comparison that isn't in the reference, (f) Omitting a comparison detailing a change from a prior study) errors to 0
+        # Note: we are NOT doing this anymore: sig_errors[-2:] = 0, 0
+
+        if matched_findings == 0:
+            return 0
+
+        if (
+            sig_present is None or matched_findings is None
+        ):  # when the template does not include the key "Clinically Significant Errors"
+            return None
+
+        return matched_findings / (matched_findings + sum(sig_errors))
+
+    def parse_error_counts(self, text, category, for_reward=False):
+
+        if category not in self.categories:
+            raise ValueError(
+                f"Category {category} is not a valid category. Please choose from {self.categories}."
+            )
+
+        # Pattern to match integers within the category, stopping at the next category or end of text
+        pattern = rf"\[{category}\]:\s*(.*?)(?:\n\s*\n|\Z)"
+        category_text = re.search(pattern, text, re.DOTALL)
+
+        # Initialize the counts
+        sum_counts = 0
+        sub_counts = [0 for i in range(6)]
+
+        # If the category is not found, return 0
+        if not category_text:
+            if for_reward:
+                # we need to know whether the category is empty or not, otherwise we overesitmate the reward
+                return None, None
+            return sum_counts, sub_counts
+        # If the category is found, but the category is empty, return 0
+        if category_text.group(1).startswith("No"):
+            return sum_counts, sub_counts
+
+        if category == "Matched Findings":
+            counts = re.findall(r"^\b\d+\b(?=\.)", category_text.group(1))
+            if len(counts) > 0:
+                sum_counts = int(counts[0])
+            return sum_counts, sub_counts
+        # Possible fine-grained error categories for categories Significant and Insignificant Clinical Errors
+        else:  # "Clinically Significant Errors" or "Clinically Insignificant Errors"
+            # Split each string at the first space and keep only the first part
+            sub_categories = [s.split(" ", 1)[0] + " " for s in self.sub_categories]
+            # Find all sub_categories in the matched text
+            matches = sorted(re.findall(r"\([a-f]\) .*", category_text.group(1)))
+
+            # this is for the gpt-4 template which assigns a number to the subcategories not letters
+            if len(matches) == 0:
+                matches = sorted(re.findall(r"\([1-6]\) .*", category_text.group(1)))
+                sub_categories = [
+                    f"({i})" + " " for i in range(1, len(self.sub_categories) + 1)
+                ]
+
+            for position, sub_category in enumerate(sub_categories):
+                # need to loop over all matches, because the sub_categories are not always in the same order
+                for match in range(len(matches)):
+                    if matches[match].startswith(sub_category):
+                        # If the sub_category is found, insert the count to sub_counts at the ordered position
+                        count = re.findall(r"(?<=: )\b\d+\b(?=\.)", matches[match])
+                        if len(count) > 0:
+                            # take the first number after the colon
+                            sub_counts[position] = int(count[0])
+            return sum(sub_counts), sub_counts
+
+    def parse_error_sentences(self, response, category):
+        """
+        Parses error sentences from a given response based of the specified category. Extracts sentences associated with each sub-categories and returns them in a dict format.
+
+        Args:
+            text (str): The input text containing error information.
+            category (str): The category to parse within the text.
+
+        Returns:
+            dict: A dictionary where keys are sub-categories and values are lists of sentences associated with those sub-categories. If the category is "Matched Findings", returns a list of sentences directly.
+        """
+        if category not in self.categories:
+            raise ValueError(
+                f"Category {category} is not a valid category. Please choose from {self.categories}."
+            )
+        pattern = rf"\[{category}\]:\s*(.*?)(?:\n\s*\n|\Z)"
+        category_text = re.search(pattern, response, re.DOTALL)
+        sub_category_dict_sentences = {}
+        for sub_category in self.sub_categories:
+            sub_category_dict_sentences[sub_category] = []
+
+        if not category_text:
+            return sub_category_dict_sentences
+        if category_text.group(1).startswith("No"):
+            return sub_category_dict_sentences
+
+        if category == "Matched Findings":
+            return (
+                category_text.group(1).rsplit(":", 1)[-1].rsplit(".", 1)[-1].split(";")
+            )
+
+        matches = sorted(re.findall(r"\([a-f]\) .*", category_text.group(1)))
+
+        if len(matches) == 0:
+            matches = sorted(re.findall(r"\([1-6]\) .*", category_text.group(1)))
+            self.sub_categories = [
+                f"({i})" + " " for i in range(1, len(self.sub_categories) + 1)
+            ]
+
+        for position, sub_category in enumerate(self.sub_categories):
+            # need to loop over all matches, because the sub_categories are not always in the same order
+            for match in range(len(matches)):
+                if matches[match].startswith(sub_category):
+                    # If the sub_category is found, add to dictionary
+                    sentences_list = (
+                        matches[match].rsplit(":", 1)[-1].split(".", 1)[-1].split(";")
+                    )
+                    sub_category_dict_sentences[self.sub_categories[position]] = (
+                        sentences_list
+                    )
+
+        return sub_category_dict_sentences
+
+    def compute_sentences(self, response):
+        # for now we only look at the significant clinical errors, which is the first category
+        return self.parse_error_sentences(response, self.categories[0])
+
+    def get_representative_sentences(self, responses):
+        list_sentences = []
+        for i in responses:
+            sentences = self.compute_sentences(i)
+            list_sentences.append(sentences)
+
+        dict_sentences = flatten_values_lists_of_list_dicts_to_dict(list_sentences)
+
+        result_sentences_dict = {}
+
+        for i in self.sub_categories:
+            sentences = dict_sentences[i]
+            sentences = [i for i in sentences if i.strip() != ""]
+            _, sentences_of_largest_cluster = compute_largest_cluster(sentences)
+            result_sentences_dict[i] = sentences_of_largest_cluster
+
+        return result_sentences_dict
+
+    def compute_accuracy(self, responses):
+        """
+        Computes the accuracy for each subcategory based on significant clinical errors and matched findings.
+
+        Args:
+            responses (list): Generated responses to evaluate.
+
+        Returns:
+            dict: accurarcies per subcategory.
+        """
+        counts = []
+        for response in responses:
+            _, sig_errors = self.parse_error_counts(response, self.categories[0])
+            counts.append(sig_errors)
+
+        counts = np.array(counts)
+
+        dict_acc = {}
+        for i in range(len(self.sub_categories)):
+            error_counts = counts[:, i]
+            # compute the accuracy for each subcategory
+            accuracy = np.mean(error_counts == 0)
+            dict_acc[self.sub_categories[i]] = accuracy
+
+        return dict_acc
+
+    def compute_summary(self):
+        """
+        Makes green summary.
+
+        Args:
+            mean_green (int): grean average.
+            mean_std (int): grean std.
+            responses (list): list of green model responses (str)
+
+        Returns:
+            str: green summary.
+        """
+        print("Computing summary ...")
+        #representative_sentences = self.get_representative_sentences(self.completions)
+        accuracies = self.compute_accuracy(self.completions)
+
+        #summary = f"\n-------------{self.model_name}----------------\n [Summary]: Green average {np.mean(self.green_scores)} and standard variation {np.std(self.green_scores)} \n [Clinically Significant Errors Analyses]: <accuracy>. <representative error>\n\n (a) False report of a finding in the candidate: {accuracies[self.sub_categories[0]]}. \n {representative_sentences[self.sub_categories[0]]} \n\n (b) Missing a finding present in the reference: {accuracies[self.sub_categories[1]]}. \n {representative_sentences[self.sub_categories[1]]} \n\n (c) Misidentification of a finding's anatomic location/position: {accuracies[self.sub_categories[2]]}. \n {representative_sentences[self.sub_categories[2]]} \n\n (d) Misassessment of the severity of a finding: {accuracies[self.sub_categories[3]]}. \n {representative_sentences[self.sub_categories[3]]} \n\n (e) Mentioning a comparison that isn't in the reference: {accuracies[self.sub_categories[4]]}. \n {representative_sentences[self.sub_categories[4]]} \n\n (f) Omitting a comparison detailing a change from a prior study: {accuracies[self.sub_categories[5]]}. {representative_sentences[self.sub_categories[5]]}.\n----------------------------------\n"
+        summary = f"\n-------------{self.model_name}----------------\n [Summary]: Green average {np.mean(self.green_scores)} and standard variation {np.std(self.green_scores)} \n [Clinically Significant Errors Analyses]: <accuracy>. <representative error>\n\n (a) False report of a finding in the candidate: {accuracies[self.sub_categories[0]]}. \n\n\n (b) Missing a finding present in the reference: {accuracies[self.sub_categories[1]]}. \n\n\n (c) Misidentification of a finding's anatomic location/position: {accuracies[self.sub_categories[2]]}. \n\n\n (d) Misassessment of the severity of a finding: {accuracies[self.sub_categories[3]]}. \n\n\n (e) Mentioning a comparison that isn't in the reference: {accuracies[self.sub_categories[4]]}. \n\n\n (f) Omitting a comparison detailing a change from a prior study: {accuracies[self.sub_categories[5]]}.\n----------------------------------\n"
+
+        print(summary)
+        return summary
+
+
+def compute(model_name, refs, hyps, output_dir="."):
+    warnings.filterwarnings("ignore", message="A decoder-only architecture is being used*") # this warning appears, despide 'padding_side='left' and correct padding
+    from sklearn.exceptions import ConvergenceWarning
+    warnings.filterwarnings("ignore", category=ConvergenceWarning, message="Number of distinct clusters.*") # test examples are copied
+    warnings.filterwarnings("ignore", category=FutureWarning, module="transformers.tokenization_utils_base")
+
+    
+    chat_template = "{% for message in messages %}\n{% if message['from'] == 'human' %}\n{{ '<|user|>\n' + message['value'] + eos_token }}\n{% elif message['from'] == 'system' %}\n{{ '<|system|>\n' + message['value'] + eos_token }}\n{% elif message['from'] == 'gpt' %}\n{{ '<|assistant|>\n'  + message['value'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
+
+    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
+        if not dist.is_initialized():
+            dist.init_process_group(
+                backend="nccl",
+            )  # 'nccl' is recommended for GPUs
+            torch.cuda.set_device(dist.get_rank())
+            if dist.get_rank() == 0:
+                print("Distributed training with", torch.cuda.device_count(), "GPUs")
+
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        trust_remote_code=False if "Phi" in model_name else True,
+        device_map={"": "cuda:{}".format(torch.cuda.current_device())},
+        torch_dtype=torch.float16,
+    )
+        
+    model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        add_eos_token=True,
+        use_fast=True,
+        trust_remote_code=True,
+        padding_side="left",
+    )
+    
+    tokenizer.chat_template = chat_template
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.clean_up_tokenization_spaces = True
+
+    inferer = Inferer(
+        dataset=[refs, hyps],
+        model=model,
+        model_name=model_name,
+        tokenizer=tokenizer,
+        output_dir=output_dir,
+        batch_size=12, # 16
+    )
+
+    t = time.time()
+
+    inferer.infer()
+
+    t = time.time() - t
+    print("Seconds per example: ", t / len(refs))
+    
+    if not is_main_process():
+        # Exit the process
+        print(f"Rank {dist.get_rank()} exiting.")
+        dist.destroy_process_group()  # Clean up the distributed processing group
+        sys.exit()  # Exit the process
+
+if __name__ == "__main__":
+    import time
+
+    refs = [
+        "Interstitial opacities without changes.",
+        "Interval development of segmental heterogeneous airspace opacities throughout the lungs . No significant pneumothorax or pleural effusion . Bilateral calcified pleural plaques are scattered throughout the lungs . The heart is not significantly enlarged .",
+        "Bibasilar atelectasis. Otherwise, no acute intrathoracic process.",
+        "Lung volumes are low, causing bronchovascular crowding. The cardiomediastinal silhouette is unremarkable. No focal consolidation, pleural effusion, or pneumothorax detected. Within the limitations of chest radiography, osseous structures are unremarkable.",
+        "Interval resolution of previously seen mild pulmonary edema with trace bilateral pleural effusions.",
+        "Lung volumes are low, causing bronchovascular crowding. The cardiomediastinal silhouette is unremarkable. No focal consolidation, pleural effusion, or pneumothorax detected. Within the limitations of chest radiography, osseous structures are unremarkable.",
+        "Bilateral pleural effusions, large on the right and small on the left. No definite focal consolidation identified, although evaluation is limited secondary to these effusions.",
+        "1. Mild left basal atelectasis. Otherwise unremarkable. 2. No definite displaced rib fracture though if there is continued concern dedicated rib series may be performed to further assess.",
+        "Interval development of segmental heterogeneous airspace opacities throughout the lungs . No significant pneumothorax or pleural effusion . Bilateral calcified pleural plaques are scattered throughout the lungs . The heart is not significantly enlarged .",
+    ]
+    hyps = [
+        "Interstitial opacities at bases without changes.",
+        "Interval resolution of previously seen mild pulmonary edema with trace bilateral pleural effusions.",
+        "Bibasilar atelectasis. Otherwise, no acute intrathoracic process.",
+        "Interval development of segmental heterogeneous airspace opacities throughout the lungs . No significant pneumothorax or pleural effusion . Bilateral calcified pleural plaques are scattered throughout the lungs . The heart is not significantly enlarged .",
+        "Endotracheal and nasogastric tubes have been removed. Changes of median sternotomy, with continued leftward displacement of the fourth inferiomost sternal wire. There is continued moderate-to-severe enlargement of the cardiac silhouette. Pulmonary aeration is slightly improved, with residual left lower lobe atelectasis. Stable central venous congestion and interstitial pulmonary edema. Small bilateral pleural effusions are unchanged.",
+        "Endotracheal and nasogastric tubes have been removed. Changes of median sternotomy, with continued leftward displacement of the fourth inferiomost sternal wire. There is continued moderate-to-severe enlargement of the cardiac silhouette. Pulmonary aeration is slightly improved, with residual left lower lobe atelectasis. Stable central venous congestion and interstitial pulmonary edema. Small bilateral pleural effusions are unchanged.",
+        "In comparison with the study of ___, the increased opacification at the right base has essentially cleared with better inspiration. Cardiac silhouette remains at the upper limits of normal in size and there is again tortuosity of the aorta without vascular congestion or pleural effusion. Biapical changes, especially on the right, are stable.",
+        "1. Mild left basal atelectasis. Otherwise unremarkable.",
+        "1. Mild left basal atelectasis. Otherwise unremarkable. 2. No definite displaced rib fracture though if there is continued concern dedicated rib series may be performed to further assess.",
+    ]
+
+    model_name = "StanfordAIMI/GREEN-radllama2-7b"
+
+    compute(model_name, refs, hyps, output_dir=".")