!pip install elasticsearch sentence_transformers transformers eland==8.12.1 # accelerate # uncomment if using GPU
!pip install datasets==2.19.2 # Remove version lock if https://github.com/huggingface/datasets/pull/6978 has been released

from huggingface_hub import notebook_login

notebook_login()

from google.colab import userdata

# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id
CLOUD_ID = userdata.get("CLOUD_ID")  # or "<YOUR CLOUD_ID>"

# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key
ELASTIC_API_KEY = userdata.get("ELASTIC_DEPL_API_KEY")  # or "<YOUR API KEY>"

from elasticsearch import Elasticsearch, helpers

# Create the client instance
client = Elasticsearch(cloud_id=CLOUD_ID, api_key=ELASTIC_API_KEY)

# Successful response!
client.info()

# Load Dataset
from datasets import load_dataset

# https://huggingface.co/datasets/MongoDB/embedded_movies
dataset = load_dataset("MongoDB/embedded_movies")

dataset

# Data Preparation

# Remove data point where plot coloumn is missing
dataset = dataset.filter(lambda x: x["fullplot"] is not None)

if "plot_embedding" in sum(dataset.column_names.values(), []):
    # Remove the plot_embedding from each data point in the dataset as we are going to create new embeddings with an open source embedding model from Hugging Face
    dataset = dataset.remove_columns("plot_embedding")

dataset["train"]

USE_ELASTICSEARCH_VECTORISATION = True

EMBEDDING_MODEL_ID = "thenlper/gte-small"
# https://huggingface.co/thenlper/gte-small's page shows the dimensions of the model
# If you use the `gte-base` or `gte-large` embedding models, the numDimension
# value in the vector search index must be set to 768 and 1024, respectively.
EMBEDDING_DIMENSIONS = 384

import locale
locale.getpreferredencoding = lambda: "UTF-8"
!(if [ "True" == $USE_ELASTICSEARCH_VECTORISATION ]; then \
  eland_import_hub_model --cloud-id $CLOUD_ID --hub-model-id $EMBEDDING_MODEL_ID --task-type text_embedding --es-api-key $ELASTIC_API_KEY --start --clear-previous; \
fi)

from sentence_transformers import SentenceTransformer

if not USE_ELASTICSEARCH_VECTORISATION:
    embedding_model = SentenceTransformer(EMBEDDING_MODEL_ID)


def get_embedding(text: str) -> list[float]:
    if USE_ELASTICSEARCH_VECTORISATION:
        raise Exception(f"Disabled when USE_ELASTICSEARCH_VECTORISATION is [{USE_ELASTICSEARCH_VECTORISATION}]")
    else:
        if not text.strip():
            print("Attempted to get embedding for empty text.")
            return []

        embedding = embedding_model.encode(text)
        return embedding.tolist()


def add_fullplot_embedding(x):
    if USE_ELASTICSEARCH_VECTORISATION:
        raise Exception(f"Disabled when USE_ELASTICSEARCH_VECTORISATION is [{USE_ELASTICSEARCH_VECTORISATION}]")
    else:
        full_plots = x["fullplot"]
        return {"embedding": [get_embedding(full_plot) for full_plot in full_plots]}


if not USE_ELASTICSEARCH_VECTORISATION:
    dataset = dataset.map(add_fullplot_embedding, batched=True)
    dataset["train"]

>>> # Needs to match the id returned from Eland
>>> # in general for Hugging Face models, you just replace the forward slash with
>>> # double underscore
>>> model_id = EMBEDDING_MODEL_ID.replace("/", "__")

>>> index_name = "movies"

>>> index_mapping = {
...     "properties": {
...         "fullplot": {"type": "text"},
...         "plot": {"type": "text"},
...         "title": {"type": "text"},
...     }
... }
>>> # define index mapping
>>> if USE_ELASTICSEARCH_VECTORISATION:
...     index_mapping["properties"]["embedding"] = {
...         "properties": {
...             "is_truncated": {"type": "boolean"},
...             "model_id": {
...                 "type": "text",
...                 "fields": {"keyword": {"type": "keyword", "ignore_above": 256}},
...             },
...             "predicted_value": {
...                 "type": "dense_vector",
...                 "dims": EMBEDDING_DIMENSIONS,
...                 "index": True,
...                 "similarity": "cosine",
...             },
...         }
...     }
>>> else:
...     index_mapping["properties"]["embedding"] = {
...         "type": "dense_vector",
...         "dims": EMBEDDING_DIMENSIONS,
...         "index": "true",
...         "similarity": "cosine",
...     }

>>> # flag to check if index has to be deleted before creating
>>> should_delete_index = True

>>> # check if we want to delete index before creating the index
>>> if should_delete_index:
...     if client.indices.exists(index=index_name):
...         print("Deleting existing %s" % index_name)
...         client.indices.delete(index=index_name, ignore=[400, 404])

>>> print("Creating index %s" % index_name)


>>> # ingest pipeline definition
>>> if USE_ELASTICSEARCH_VECTORISATION:
...     pipeline_id = "vectorize_fullplots"

...     client.ingest.put_pipeline(
...         id=pipeline_id,
...         processors=[
...             {
...                 "inference": {
...                     "model_id": model_id,
...                     "target_field": "embedding",
...                     "field_map": {"fullplot": "text_field"},
...                 }
...             }
...         ],
...     )

...     index_settings = {
...         "index": {
...             "default_pipeline": pipeline_id,
...         }
...     }
>>> else:
...     index_settings = {}

>>> client.options(ignore_status=[400, 404]).indices.create(
...     index=index_name, mappings=index_mapping, settings=index_settings
... )

>>> from elasticsearch.helpers import BulkIndexError


>>> def batch_to_bulk_actions(batch):
...     for record in batch:
...         action = {
...             "_index": "movies",
...             "_source": {
...                 "title": record["title"],
...                 "fullplot": record["fullplot"],
...                 "plot": record["plot"],
...             },
...         }
...         if not USE_ELASTICSEARCH_VECTORISATION:
...             action["_source"]["embedding"] = record["embedding"]
...         yield action


>>> def bulk_index(ds):
...     start = 0
...     end = len(ds)
...     batch_size = 100
...     if USE_ELASTICSEARCH_VECTORISATION:
...         # If using auto-embedding, bulk requests can take a lot longer,
...         # so pass a longer request_timeout here (defaults to 10s), otherwise
...         # we could get Connection timeouts
...         batch_client = client.options(request_timeout=600)
...     else:
...         batch_client = client
...     for batch_start in range(start, end, batch_size):
...         batch_end = min(batch_start + batch_size, end)
...         print(f"batch: start [{batch_start}], end [{batch_end}]")
...         batch = ds.select(range(batch_start, batch_end))
...         actions = batch_to_bulk_actions(batch)
...         helpers.bulk(batch_client, actions)


>>> try:
...     bulk_index(dataset["train"])
>>> except BulkIndexError as e:
...     print(f"{e.errors}")

>>> print("Data ingestion into Elasticsearch complete!")

def vector_search(plot_query):
    if USE_ELASTICSEARCH_VECTORISATION:
        knn = {
            "field": "embedding.predicted_value",
            "k": 10,
            "query_vector_builder": {
                "text_embedding": {
                    "model_id": model_id,
                    "model_text": plot_query,
                }
            },
            "num_candidates": 150,
        }
    else:
        question_embedding = get_embedding(plot_query)
        knn = {
            "field": "embedding",
            "query_vector": question_embedding,
            "k": 10,
            "num_candidates": 150,
        }

    response = client.search(index="movies", knn=knn, size=5)
    results = []
    for hit in response["hits"]["hits"]:
        id = hit["_id"]
        score = hit["_score"]
        title = hit["_source"]["title"]
        plot = hit["_source"]["plot"]
        fullplot = hit["_source"]["fullplot"]
        result = {
            "id": id,
            "_score": score,
            "title": title,
            "plot": plot,
            "fullplot": fullplot,
        }
        results.append(result)
    return results


def pretty_search(query):

    get_knowledge = vector_search(query)

    search_result = ""
    for result in get_knowledge:
        search_result += f"Title: {result.get('title', 'N/A')}, Plot: {result.get('fullplot', 'N/A')}\n"

    return search_result

>>> # Conduct query with retrival of sources, combining results into something that
>>> # we can feed to Gemma
>>> def combined_query(query):
...     source_information = pretty_search(query)
...     return f"Query: {query}\nContinue to answer the query by using these Search Results:\n{source_information}."


>>> query = "What is the best romantic movie to watch and why?"
>>> combined_results = combined_query(query)

>>> print(combined_results)

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
# CPU Enabled uncomment below 👇🏽
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
# GPU Enabled use below 👇🏽
# model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", device_map="auto")

>>> def rag_query(query):
...     combined_information = combined_query(query)

...     # Moving tensors to GPU
...     input_ids = tokenizer(combined_information, return_tensors="pt")  # .to("cuda") # Add if using GPU
...     response = model.generate(**input_ids, max_new_tokens=700)

...     return tokenizer.decode(response[0], skip_special_tokens=True)


>>> print(rag_query("What's a romantic movie that I can watch with my wife?"))