updated file structure

README.md

@@ -10,4 +10,3 @@ pinned: false
 license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,28 @@
+import gradio as gr
+import random
+import time
+
+from common import DATA
+from config import DEFAULT_BOT_MESSAGE
+from inference_hf import rag_chatbot
+
+
+DATA = DATA.add_faiss_index("embedding")
+
+
+async def predict(message, chat_history):
+    bot_message = rag_chatbot(message, k=3)
+    chat_history.append((message, bot_message))
+    return "", chat_history
+
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    chatbot = gr.Chatbot(
+        value=[[None, DEFAULT_BOT_MESSAGE]], label="ReformaPensional-Llama3"
+    )
+    msg = gr.Textbox(placeholder="Haz aquí tu pregunta")
+    clear = gr.ClearButton([msg, chatbot])
+    msg.submit(predict, [msg, chatbot], [msg, chatbot])
+
+if __name__ == "__main__":
+    demo.launch()

common.py

@@ -0,0 +1,79 @@
+import torch
+from datasets import Dataset as hfd
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    pipeline,
+)
+
+from config import DATASET_HF_NAME, LLAMA3_CHECKPOINT
+
+# Adapted from HF https://huggingface.co/blog/not-lain/rag-chatbot-using-llama3
+
+
+def search_topk(
+    data: hfd,
+    feature_extractor: SentenceTransformer,
+    query: str,
+    k: int = 3,
+    embedding_col: str = "embedding",
+):
+    """a function that embeds a new query and returns the most probable results"""
+    embedded_query = feature_extractor.encode(query)  # embed new query
+    scores, retrieved_examples = data.get_nearest_examples(  # retrieve results
+        embedding_col,
+        embedded_query,  # compare our new embedded query with the dataset embeddings
+        k=k,  # get only top k results
+    )
+    return scores, retrieved_examples
+
+
+def format_prompt(
+    prompt: str, retrieved_documents: hfd, k: int, text_col: str = "chunk"
+):
+    """using the retrieved documents we will prompt the model to generate our responses"""
+    PROMPT = f"Question:{prompt}\nContext:"
+    for idx in range(k):
+        PROMPT += f"{retrieved_documents[text_col][idx]}\n"
+    return PROMPT
+
+
+# Quantization Config
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+# Tokenizer & Model
+# You must request access to the checkpoints
+TOKENIZER = AutoTokenizer.from_pretrained(LLAMA3_CHECKPOINT)
+MODEL = AutoModelForCausalLM.from_pretrained(
+    LLAMA3_CHECKPOINT,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    quantization_config=bnb_config,
+)
+TERMINATORS = [TOKENIZER.eos_token_id, TOKENIZER.convert_tokens_to_ids("<|eot_id|>")]
+
+DATA = load_dataset(DATASET_HF_NAME)["train"]
+
+TEXT_GENERATION_PIPELINE = pipeline(
+    model=MODEL,
+    tokenizer=TOKENIZER,
+    task="text-generation",
+    device_map="auto",
+)
+TEXT_GENERATION_PIPELINE.tokenizer
+
+PIPELINE_INFERENCE_ARGS = {
+    "max_new_tokens": 256,
+    "eos_token_id": TERMINATORS,
+    "do_sample": True,
+    "temperature": 0.1,
+    "top_p": 0.9,
+}

config.py

@@ -0,0 +1,13 @@
+FEATURE_EXTRACTOR_CHECKPOINT = "BAAI/bge-large-en-v1.5"
+DATASET_HF_NAME = "jhonparra18/reforma-pensional-col"
+LLAMA3_CHECKPOINT = "meta-llama/Meta-Llama-3-8B-Instruct"
+SYS_PROMPT_HF = """
+    Eres un asistente automático que brinda información referente a la reforma
+    pensional del actual gobierno, tu meta es responder a las preguntas y cuestionamientos
+    en la manera más precisa y haciendo referencia a los textos de la reforma.
+    Siempre responde respecto a la información que se proporciona.
+    Tu respuesta jamás debe corresponder a cosas por fuera del texto que se te da.
+    """
+MAX_TOKENS_INPUT = 2000
+
+DEFAULT_BOT_MESSAGE = "Hola! Soy un chatbot construido con LLama3 para responder preguntas de la reforma pensional. Haz las preguntas que desees"

inference_hf.py

@@ -0,0 +1,47 @@
+from common import DATA, MODEL, TERMINATORS, TOKENIZER, format_prompt, search_topk
+from config import MAX_TOKENS_INPUT, SYS_PROMPT_HF
+from preprocessing import FEATURE_EXTRACTOR
+
+
+def generate(formatted_prompt):
+    formatted_prompt = formatted_prompt[:MAX_TOKENS_INPUT]  # to avoid GPU OOM
+    messages = [
+        {"role": "system", "content": SYS_PROMPT_HF},
+        {"role": "user", "content": formatted_prompt},
+    ]
+
+    input_ids = TOKENIZER.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt"
+    ).to(MODEL.device)
+    outputs = MODEL.generate(
+        input_ids,
+        max_new_tokens=512,
+        eos_token_id=TERMINATORS,
+        do_sample=True,
+        temperature=0.1,
+        top_p=0.9,
+    )
+    response = outputs[0]
+    return TOKENIZER.decode(response[input_ids.shape[-1] :], skip_special_tokens=True)
+
+
+def rag_chatbot(prompt: str, k: int = 2, return_user: bool = False):
+    _, retrieved_documents = search_topk(
+        DATA, FEATURE_EXTRACTOR, prompt, k, embedding_col="embedding"
+    )
+    formatted_prompt = format_prompt(prompt, retrieved_documents, k, text_col="chunk")
+    bot_response = generate(formatted_prompt)
+    return (
+        f"[USER]: {prompt}\n\n[ASSISTANT]: {bot_response}"
+        if return_user
+        else bot_response
+    )
+
+
+if __name__ == "__main__":
+    # example RAG Pipeline using HuggingFace
+    DATA = DATA.add_faiss_index("embedding")
+    prompt = """indicame qué va a pasar en la reforma pensional con los fondos en el pilar
+    contributivo de prima media, podré pedir el dinero de vuelta cuando tenga la edad si no
+    cumplo con las semanas cotizadas?"""
+    print(rag_chatbot(prompt, k=3, return_user=True))

preprocessing.py

@@ -0,0 +1,35 @@
+import textract
+from datasets import Dataset as hfd
+from sentence_transformers import SentenceTransformer
+
+from config import FEATURE_EXTRACTOR_CHECKPOINT
+
+FEATURE_EXTRACTOR = SentenceTransformer(FEATURE_EXTRACTOR_CHECKPOINT)
+
+
+def encode_sentence(instance: hfd, text_col: str):
+    return {
+        "embedding": FEATURE_EXTRACTOR.encode(
+            instance[text_col], normalize_embeddings=True
+        )
+    }
+
+
+def parse_pdf(pdf_path: str):
+    """Gets text from a pdf file using textract"""
+    txt = textract.process(pdf_path, method="pdfminer", encoding="latin-1").decode()
+    return txt
+
+
+def chunk_text(text: str, split_sentence="ARTÍCULO"):
+    """creates chunks of texts using a split_sentence"""
+    chunks = [
+        {"chunk": split_sentence + " " + c.replace("\n", " ").strip()}
+        for c in text.split(split_sentence)
+    ]
+    return chunks
+
+
+def create_df(text_chunks: list[dict[str]]):
+    "creates a HuggingFace dataset based on a list of dicts [str,str]"
+    return hfd.from_list(text_chunks)

requirements.txt

@@ -0,0 +1,6 @@
+transformers
+datasets
+sentence-transformers
+faiss-cpu
+accelerate
+bitsandbytes