!pip install -q -U distilabel "farm-haystack[preprocessing]"
!pip install -q -U "distilabel[hf-inference-endpoints, argilla]"

import os
from typing import Dict

from distilabel.llm import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline, pipeline
from distilabel.tasks import TextGenerationTask, SelfInstructTask, Prompt

from datasets import Dataset
from haystack.nodes import PDFToTextConverter, PreProcessor

os.environ["HF_TOKEN"] = ""
os.environ["HF_INFERENCE_ENDPOINT_NAME"] = "aws-notus-7b-v1-3184"
os.environ["HF_NAMESPACE"] = "argilla"
os.environ["OPENAI_API_KEY"] = ""

class QuestionAnsweringTask(TextGenerationTask):
    def generate_prompt(self, question: str) -> str:
        return Prompt(
            system_prompt=self.system_prompt,
            formatted_prompt=question,
        ).format_as(
            "llama2"
        )  # type: ignore

    def parse_output(self, output: str) -> Dict[str, str]:
        return {"answer": output.strip()}

    @property
    def input_args_names(self) -> list[str]:
        return ["question"]

    @property
    def output_args_names(self) -> list[str]:
        return ["answer"]

llm = InferenceEndpointsLLM(
    endpoint_name_or_model_id=os.getenv("HF_INFERENCE_ENDPOINT_NAME"),  # type: ignore
    endpoint_namespace=os.getenv("HF_NAMESPACE"),  # type: ignore
    token=os.getenv("HF_TOKEN") or None,
    task=QuestionAnsweringTask(),
)

generation = llm.generate([{"question": "What's the second most populated city in Denmark?"}])
generation[0][0]["parsed_output"]["answer"]

%%bash

if [ ! -f "The-AI-Act.pdf" ]; then
    wget -q https://artificialintelligenceact.eu/wp-content/uploads/2021/08/The-AI-Act.pdf
fi

# The converter turns the PDF into text we can process easily
converter = PDFToTextConverter(remove_numeric_tables=True, valid_languages=["en"])

# Preprocessing pipelines can have several steps.
# Ours clean empty lines, header, footers and whitespaces
# and split the text into 150-char long batches, respecting
# where the sentences naturally end and begin.
preprocessor = PreProcessor(
    clean_empty_lines=True,
    clean_whitespace=True,
    clean_header_footer=True,
    split_by="word",
    split_length=150,
    split_respect_sentence_boundary=True,
)

doc = converter.convert(file_path="The-AI-Act.pdf", meta=None)[0]
docs = preprocessor.process([doc])
print(f"Documents: 1\nBatches: {len(docs)}")

inputs = [doc.content for doc in docs]
inputs[0][0:500]

instructions_dataset = Dataset.from_dict({"input": inputs[0:50]})

instructions_dataset

instructions_task = SelfInstructTask(
    application_description="A assistant that can answer questions about the AI Act made by the European Union."
)

instructions_generator = InferenceEndpointsLLM(
    endpoint_name_or_model_id=os.getenv("HF_INFERENCE_ENDPOINT_NAME"),  # type: ignore
    endpoint_namespace=os.getenv("HF_NAMESPACE"),  # type: ignore
    token=os.getenv("HF_TOKEN") or None,
    task=instructions_task,
)

instructions_pipeline = Pipeline(generator=instructions_generator)

generated_instructions = instructions_pipeline.generate(dataset=instructions_dataset, num_generations=1, batch_size=8)

>>> instructions = []
>>> for generations in generated_instructions["instructions"]:
...     for generation in generations:
...         instructions.extend(generation)

>>> print(f"Number of generated instructions: {len(instructions)}")

>>> for instruction in instructions[:5]:
...     print(instruction)

generated_instructions[0]

instructions_rg_dataset = generated_instructions.to_argilla()
instructions_rg_dataset[0]

instructions_rg_dataset.push_to_argilla(name=f"notus_AI_instructions")

preference_pipeline = pipeline(
    "preference",
    "instruction-following",
    generator=InferenceEndpointsLLM(
        endpoint_name_or_model_id=os.getenv("HF_INFERENCE_ENDPOINT_NAME"),  # type: ignore
        endpoint_namespace=os.getenv("HF_NAMESPACE", None),
        task=TextGenerationTask(),
        max_new_tokens=256,
        num_threads=2,
        temperature=0.3,
    ),
    max_new_tokens=256,
    num_threads=2,
    api_key=os.getenv("OPENAI_API_KEY", None),
    temperature=0.0,
)

remote_dataset = rg.FeedbackDataset.from_argilla("notus_AI_instructions", workspace="admin")
instructions_dataset = remote_dataset.pull(max_records=100)  # get first 100 records

instructions_dataset = instructions_dataset.format_as("datasets")
instructions_dataset

instructions_dataset[0]

instructions_dataset = instructions_dataset.rename_columns({"input": "context", "instruction": "input"})

preference_dataset = preference_pipeline.generate(
    instructions_dataset,  # type: ignore
    num_generations=2,
    batch_size=8,
    display_progress_bar=True,
)

preference_dataset[0]

import argilla as rg

# Replace api_url with the url to your HF Spaces URL if using Spaces
# Replace api_key if you configured a custom API key
rg.init(api_url="http://localhost:6900", api_key="owner.apikey", workspace="admin")

# Uploading the Preference Dataset
preference_rg_dataset = preference_dataset.to_argilla()

# Adding the context as a metadata property in the new Feedback dataset, as this
# information will be useful later.
for record_feedback, record_huggingface in zip(preference_rg_dataset, preference_dataset):
    record_feedback.metadata["context"] = record_huggingface["context"]

preference_rg_dataset.push_to_argilla(name=f"notus_AI_preference")