from transformers import AutoModelForCausalLM, AutoTokenizer, OlmoForCausalLM
import torch
import csv 
import pandas as pd 

#load in the different models 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.cuda.get_device_properties(0))
olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-1124-7B").to(device)
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-2-1124-7B")

#priming prompt
prompt_1 = "For the GIVEN DATA, Please categorize it based on the following numbered characteristics: \n\n 1: YES/NO (Characteristic 1. This is an English language empirical study. English language empirical studies discuss data or observations.) \n 2: YES/NO (Characteristic 2. This focuses on free and open source software (FOSS). The focus of this paper is on free or open source software projects and ecosystems.) \n 3: YES/NO (Characteristic 3. This focuses on FOSS project evolution. FOSS project evolution describes changes to free and open source projects and ecosystems.) \n 4: YES/NO (Characteristic 4. This focuses on FOSS project adaptation. FOSS project adaptation describes the intentional changes made by projects to better align with the project's broader environment.) \n\n Only respond with the appropriate number followed by 'YES' if the characteristic is present in the provided data or 'NO' if it is not (e.g. '1. YES; 2. NO;'). Do not provide any additional information."

example_4 = "Example 4: TITLE - Analysis of Open Source Software Evolution Using Evolution Curve Method \n ABSTRACT - Design and evolution of modem information systems is influenced by many   factors: technical, organizational, social, and psychological. This is   especially true for open source software systems (OSSS), when many   developers from different backgrounds interact, share their ideas and   contribute towards the development and improvement of a software   product. The evolution of all OSSS is a continuous process of source   code development, adaptation, improvement and maintenance. Studying   changes to the various characteristics of source code can help us   understand the evolution of a software system. In this paper, the   software evolution process is analyzed using a proposed Evolution curve   (E-curve) method, which is based on information theoretic metrics of   source code. The method allows identifying major evolution stages and   transition points of an analyzed software system. The application of the   E-curves is demonstrated for the eMule system. .\n CATEGORIES: 1. YES; 2. YES; 3.YES; 4. NO"

example_1 = "Example 1: TITLE - Thermal Insulation Properties of Milkweed Floss Nonwovens: Influence of   Temperature, Relative Humidity, and Fiber Content \n ABSTRACT - This study investigated the influence of fiber content, temperature, and   relative humidity on the thermal insulation properties of nonwoven mats   made of seed fibers from Asclepias Syriaca, commonly known as milkweed   floss. Nonwoven mats with a 1-inch thickness were produced by uniformly   arranging milkweed fibers within a mold. Various quantities of fiber   were employed to obtain nonwoven mats with a fiber content ranging from   5 to 35 kg/m3. Thermal conductivity and thermal diffusivity were   measured across diverse relative humidity levels and temperatures.   Simultaneously, milkweed floss samples were exposed to identical   environmental conditions to assess the moisture regain and specific heat   capacities of the fiber. The specific heat capacity of milkweed and   thermal conductivity of the nonwovens exhibited a linear increase with   temperature. The thermal diffusivity and thermal conductivity of the   nonwovens decreased with rising fiber content. The thermal insulation   properties of the nonwovens remained partially stable below 30\\%   relative humidity but substantially deteriorated at higher levels. The   nonwovens exhibited optimal thermal insulation properties at a fiber   content between 20 and 25 kg/m3. The results of this study highlighted   several technical advantages of employing milkweed floss as a   sustainable and lightweight solution for thermal insulation. \n CATEGORIES: 1. YES; 2. NO; 3. NO; 4. NO;"

example_3 = "Example 3: TITLE - Social network structures in open source software development teams \n ABSTRACT - Drawing on social network theories and previous studies, this research   examines the dynamics of social network structures in open source   software (OSS) teams. Three projects were selected from SourceForge.net   in terms of their similarities as well as their differences. Monthly   data were extracted from the bug tracking systems in order to achieve a   longitudinal view of the interaction pattern of each project. Social   network analysis was used to generate the indices of social structure.   The finding suggests that the interaction pattern of OSS projects   evolves from a single hub at the beginning to a corel periphery model as   the projects move forward.\n CATEGORIES: 1. YES; 2. YES; 3. NO; 4. NO"

example_2 = "Example 2: TITLE - An Exploratory Mixed-methods Study on General Data Protection Regulation (GDPR) Compliance in Open-Source Software \n ABSTRACT- Background: Governments worldwide are considering data privacy regulations. These laws, such as the European Union’s General Data Protection Regulation (GDPR), require software developers to meet privacy-related requirements when interacting with users’ data. Prior research describes the impact of such laws on software development, but only for commercial software. Although open-source software is commonly integrated into regulated software, and thus must be engineered or adapted for compliance, we do not know how such laws impact open-source software development. Aims: To understand how data privacy laws affect open-source software (OSS) development, we focus on the European Union’s GDPR, as it is the most prominent such law. We investigated how GDPR compliance activities influence OSS developer activity (RQ1), how OSS developers perceive fulfilling GDPR requirements (RQ2), the most challenging GDPR requirements to implement (RQ3), and how OSS developers assess GDPR compliance (RQ4). Method: We distributed an online survey to explore perceptions of GDPR implementations from open-source developers (N=56). To augment this analysis, we further conducted a repository mining study to analyze development metrics on pull requests (N=31,462) submitted to open-source GitHub repositories. Results: Our results suggest GDPR policies complicate OSS development and introduce challenges, primarily regarding the management of users’ data, implementation costs and time, and assessments of compliance. Moreover, we observed negative perceptions of the GDPR from OSS developers and significant increases in development activity, in particular metrics related to coding and reviewing, on GitHub pull requests related to GDPR compliance. Conclusions: Our findings provide future research directions and implications for improving data privacy policies, motivating the need for relevant resources and automated tools to support data privacy regulation implementation and compliance efforts in OSS. \n CATEGORIES: 1. YES; 2. YES; 3. YES; 4. YES;"

with open("cites/053025_man_filtered_dedup.csv", mode='r', newline='') as file:
    reader = csv.reader(file)
    array_of_categorizations = []
    index = -1
    for row in reader:
        index += 1
        if index <= 0:
            continue
        cite_dict = {}
        #organizing the data from each citation
        cite_dict['key'] = row[0]
        cite_dict['title'] = row[1]
        cite_dict['abstract'] = row[2]
        #prompt construction
        given_data = f"GIVEN DATA: Title - {cite_dict['title']} \n Abstract - {cite_dict['abstract']}"
        prompt = f"{prompt_1}\n\n{example_1}\n\n{example_2}\n\n{example_3}\n\n{example_4}\n\n{given_data}\n"
        #handoff to the model
        inputs = tokenizer(prompt, return_tensors='pt', return_token_type_ids=False).to(device)
        #deterministic sampling and getting the response back 
        response = olmo.generate(**inputs, max_new_tokens=256, do_sample=False)
        response_txt = tokenizer.batch_decode(response, skip_special_tokens=True)[0]
        #getting the resulting codes 
        codes_id = response_txt.rfind("CATEGORIES:")
        if codes_id != -1:
            result = response_txt[codes_id + len("CATEGORIES:"):].strip()
        else:
            cite_dict["1"] = "NULL"
            cite_dict["2"] = "NULL"
            cite_dict["3"] = "NULL"
            cite_dict["4"] = "NULL"
        #writing them to the citation_dict
        for item in result.strip(";").split(";"):
            key_value = item.strip().split('. ')
            if len(key_value) == 2:
                key = key_value[0]
                value = key_value[1]
                cite_dict[key] = value
        array_of_categorizations.append(cite_dict)
    #CSV everything
    df = pd.DataFrame(array_of_categorizations)
    df.to_csv('060225_olmo_categorized_citations.csv', index=False)
