from distilabel.embeddings import SentenceTransformerEmbeddings
from distilabel.llms import  InferenceEndpointsLLM
from distilabel.mixins.runtime_parameters import RuntimeParameter
from distilabel.pipeline import Pipeline
from distilabel.steps import EmbeddingGeneration, ConversationTemplate, FaissNearestNeighbour, StepInput, StepOutput, step
from distilabel.steps.tasks import MagpieGenerator, TextGeneration


REASONING_PROMPT = """
"You are an AI assistant specialized in logical thinking and problem-solving.
Your purpose is to help users work through complex ideas, analyze situations, and draw conclusions based on given information.
Approach each query with structured thinking, break down problems into manageable parts, and guide users through the reasoning process step-by-step."
"""

REFLECTION_PROMPT = """
You're an AI assistant that responds to the user with maximum accuracy.
To do so, your first task is to think what the user is asking for, thinking step by step.
During this thinking phase, you will have reflections that will help you clarifying ambiguities.
In each reflection you will list the possibilities and finally choose one. Between reflections, you can think again.
At the end of the thinking, you must draw a conclusion.
You only need to generate the minimum text that will help you generating a better output, don't be verbose while thinking.
Finally, you will generate an output based on the previous thinking.

This is the output format you have to follow:

```
<thinking>
Here you will think about what the user asked for.
</thinking>

<reflection>
This is a reflection.
</reflection>

<reflection>
This is another reflection.
</reflection>

</thinking>

<output>
Here you will include the output.
</output>
```
""".lstrip()

@step(inputs=["instruction", "generation"], step_type="global")
def FilteringStep(inputs: StepInput) -> StepOutput:
    yield [
        input
        for input in inputs
        # if instruction and generation are not null or empty
        if (input["instruction"] is not None and input["instruction"] != '') and
           (input["generation"] is not None and input["generation"] != '')
    ]

with Pipeline("Reflection Dataset") as pipeline:
    magpie = MagpieGenerator(
        only_instruction=True,
        num_rows=20000,
        batch_size=5,
        system_prompt=REASONING_PROMPT,
        llm=InferenceEndpointsLLM(
            model_id="meta-llama/Meta-Llama-3.1-70B-Instruct",
            tokenizer_id="meta-llama/Meta-Llama-3.1-70B-Instruct",
            magpie_pre_query_template="llama3",
            generation_kwargs={"temperature": 0.8, "max_new_tokens": 2048},
        )
    )

    response = TextGeneration(
        input_batch_size=5,
        system_prompt=REFLECTION_PROMPT,
        llm=InferenceEndpointsLLM(
            model_id="meta-llama/Meta-Llama-3.1-70B-Instruct",
            tokenizer_id="meta-llama/Meta-Llama-3.1-70B-Instruct",
            generation_kwargs={"temperature": 0.8, "max_new_tokens": 2048},
        ),
    )

    conversation = ConversationTemplate(
        input_mappings={
            "instruction": "instruction",
            "response": "generation",
        },
        output_mappings={
            "conversation": "text",
        },
    )

    filter_null = FilteringStep()

    embedding_generation = EmbeddingGeneration(
        embeddings=SentenceTransformerEmbeddings(
            model="sentence-transformers/all-MiniLM-L6-v2",
            trust_remote_code=True,
            device="cpu",
        ),
    )

    faiss = FaissNearestNeighbour()

    magpie >> response >> conversation >> filter_null >> embedding_generation >> faiss

if __name__ == "__main__":
    distiset = pipeline.run(
        use_cache=False,
    )

    distiset.push_to_hub(
        repo_id="thesven/Reflective-MAGLLAMA-v0.1",
        private=True,
        include_script=True
    )