from transformers import T5Tokenizer, T5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large").to('cuda') # With GPU --> .to('cuda')

"It is expressly understood that Commerce may unilaterally cancel this Contract for Contractor’s refusal to comply with this provision."

input_prompt = f"""
    Identify the clause type. If it is none, reply None

    Text: "{clause_text}"

    Contractual clause types:
    1. Termination condition
    2. Performance obligation
    3. Payment obligation
    4. Compliance obligation
    5. Support or warranty obligation
    6. Insurance obligation
    7. Quality obligation
    8. Audit obligation
    9. Contract renewal
    10. Exceptions or exemptions
    11. Indemnification
    12. Jurisdiction
    """

input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids.to("cuda")
outputs = model.generate(input_ids, max_length=512)
category = tokenizer.decode(outputs[0])

>>> print(category)
"<pad> 1. </s>"

import json
import pandas as pd
import numpy as np
import torch

from datasets import load_dataset, load_metric, DatasetDict

from transformers import AutoTokenizer, T5Tokenizer

zs = load_dataset('csv', data_files='zeroshot.csv', split="train")
# 90% train, 10% test + validation
train_testvalid = zs.train_test_split(test_size=0.1)
# Split the 10% test + valid in half test, half valid
test_valid = train_testvalid['test'].train_test_split(test_size=0.5)
# gather everyone if you want to have a single DatasetDict
zs = DatasetDict({
    'train': train_testvalid['train'],
    'test': test_valid['test'],
    'valid': test_valid['train']})



model_checkpoint = 'google/flan-t5-large'
tokenizer = T5Tokenizer.from_pretrained(model_checkpoint)
max_input_length = 512 # change to setting later
max_target_length = 512 # change to setting later


def preprocess_data(examples):

  inputs = []

  for clause_text in examples['clause']:
    
    input_prompt = f"""
        Identify the clause type. If it is none, reply None.

        Text: "{clause_text}"

        Contractual clause types:
        1. Termination condition
        2. Performance obligation
        3. Payment obligation
        4. Compliance obligation
        5. Support or warranty obligation
        6. Insurance obligation
        7. Quality obligation
        8. Audit obligation
        9. Contract renewal
        10. Exceptions or exemptions
        11. Indemnification
        12. Jurisdiction
        """
    
    inputs.append(input_prompt)

  model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)

  # Setup the tokenizer for targets
  with tokenizer.as_target_tokenizer():
    labels = tokenizer(examples["flan_fewshot_cat"], max_length=max_target_length, 
                       truncation=True)

  model_inputs["labels"] = labels["input_ids"]
  return model_inputs

tokenized_datasets = zs.map(preprocess_data, batched=True)

from transformers import AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer

model_name = "flan-t5-large-skim"
model_dir = f"bigdata/{model_name}"

args = Seq2SeqTrainingArguments(
    model_dir,
    evaluation_strategy="steps",
    eval_steps=100,
    logging_strategy="steps",
    logging_steps=100,
    save_strategy="steps",
    save_steps=200,
    learning_rate=4e-5,
    per_device_train_batch_size=1, 
    gradient_accumulation_steps=4,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=1,
    predict_with_generate=True,
    fp16=False,
    load_best_model_at_end=True,
    metric_for_best_model="rouge1",
)

data_collator = DataCollatorForSeq2Seq(tokenizer)


# Function that returns an untrained model to be trained
from transformers import T5Tokenizer, T5ForConditionalGeneration
def model_init():
    return T5ForConditionalGeneration.from_pretrained(model_checkpoint)


trainer = Seq2SeqTrainer(
    model=model_init(),
    args=args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["valid"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

trainer.train()

trainer.save_model("bigdata/flan-t5-large-skim") # save model locally
trainer.push_to_hub() # push to hugging face hub

# Extract text from the .docx file
doc = Document(docx_stream)
clause_text = ""

for paragraph in doc.paragraphs:
    clause_text += paragraph.text + "\n"

import spacy

from spacy.cli import download

# Specify the model name
model_name = 'en_core_web_sm'

# Check if the model is already installed
if model_name not in spacy.util.get_installed_models():
    # If not installed, download the model
    download(model_name)

nlp = spacy.load(model_name)

def split_into_sentences(text):
    sentences = list(nlp(text).sents)
    sentences = [sentence.text for sentence in sentences]
    return sentences

>>> text = r"Contractor shall allow public access to all documents, papers, letters or other materials made or received by Contractor in conjunction with this Contract, unless the records are exempt from section 24(a) of Article I of the State Constitution and section 119.07(1), F.S.  It is expressly understood that Commerce may unilaterally cancel this Contract for Contractor’s refusal to comply with this provision."
>>> [text.text for text in list(nlp(text).sents)]

['Contractor shall allow public access to all documents, papers, letters or other materials made or received by Contractor in conjunction with this Contract, unless the records are exempt from section 24(a) of Article I of the State Constitution and section 119.07(1), F.S.  ', 'It is expressly understood that Commerce may unilaterally cancel this Contract for Contractor’s refusal to comply with this provision.']

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_distances
from itertools import combinations


def remove_repetitive_strings(strings, similarity_threshold=80):
    # Use TF-IDF Vectorizer to embed each line as a vector
    vectorizer = TfidfVectorizer()
    vectors = vectorizer.fit_transform(strings)

    # Calculate pairwise cosine distances
    distances = cosine_distances(vectors)

    # Identify pairs with similarity below the threshold
    filtered_pairs = [(strings[i], strings[j]) for i, j in combinations(range(len(strings)), 2)
                      if distances[i, j] <= (1 - similarity_threshold / 100)]

    # Create a set of strings to keep track of unique strings
    unique_strings = set()

    # Add strings from filtered pairs to unique_strings
    for string1, string2 in filtered_pairs:
        unique_strings.add(string1)
        unique_strings.add(string2)

    # Filter the original list based on unique strings
    filtered_strings = [string for string in strings if string not in unique_strings]

    return filtered_strings

# Flask side:
@app.route('/api/generate_clause', methods=['POST'])
def generate_clause():

    def generate():
        progress=0
        clause_text = request.json['text']

        lotexts = splitter(clause_text) # splits text according to preprocessing rules

        """
        Takes in a list of texts (lotexts) that has already been processed. could be a word doc file or just simple text, either way, they just need to be in an orderly list of texts.
        """
        lodict = {}
        for i,text in enumerate(lotexts):
            progress = (i+1)/len(lotexts)

            input_prompt = f"""
            Identify the clause type. If it is none, reply None

            Text: "{text}"

            Contractual clause types:
            1. Termination condition
            2. Performance obligation
            3. Payment obligation
            4. Compliance obligation
            5. Support or warranty obligation
            6. Insurance obligation
            7. Quality obligation
            8. Audit obligation
            9. Contract renewal
            10. Exceptions or exemptions
            11. Indemnification
            12. Jurisdiction
            """

            # error handling for model initialization
            output = query({
                "inputs": input_prompt, "wait_for_model": True
            })

            category = output[0]['generated_text']

            lodict.update({text: category})

            json_data = {'text_block': clause_text, 'lookup': lodict, 'progress': progress}
            json_string = json.dumps(json_data)

            yield json_string  # Add a special object to separate JSON streamed objects

    # Use Response and stream_with_context to handle streaming
    return Response(stream_with_context(generate()),
                    content_type='text/event-stream')