Add functions to implement missing features

functions.py

@@ -1,10 +1,11 @@
 import os
-
 import requests
+import random
 import torch
 from bs4 import BeautifulSoup
 from peft import PeftConfig, PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, AutoModel
+from datasets import DatasetDict
 
 # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
@@ -12,6 +13,49 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 generation_config = GenerationConfig(temperature=.8,
                                      top_p=0.75,
                                      top_k=40)
+device = 'cuda'
+
+shared = {
+    'answer_context': None,
+    'embeddings_dataset': None
+}
+
+
+def get_nearest_examples(question: str, k: int):
+    print(['get_nearest_examples', 'start'])
+    question_embedding = get_embeddings([question]).cpu().detach().numpy()
+    embeddings_dataset = shared['embeddings_dataset']
+    scores, samples = embeddings_dataset.get_nearest_examples(
+        "embeddings", question_embedding, k)
+    print(['get_nearest_examples', 'scores and samples'])
+    for i in range(len(scores)):
+        print([scores[i], samples[i]])
+    print(['get_nearest_examples', 'end'])
+    return samples
+
+
+def get_embeddings(text):
+    print(['get_embeddings', 'start'])
+    encoded_input = tokenizer(
+        text, padding=True, truncation=True, return_tensors="pt")
+    encoded_input = {k: v.to('cuda') for k, v in encoded_input.items()}
+    model_output = model(**encoded_input)
+    model_output = model_output.last_hidden_state[:, 0]
+    emb_item = model_output.detach().cpu().numpy()[0]
+    print(['get_embeddings', 'end'])
+    return emb_item
+
+
+def build_faiss_index(text):
+    print(['build_faiss_index', 'start'])
+    text_list = split_text(text)
+    emb_list = []
+    for item in text_list:
+        emb_list.append({"embeddings": get_embeddings(item)})
+    dataset = DatasetDict({'train': emb_list})
+    dataset.add_faiss_index(column="embeddings")
+    shared['embeddings_dataset'] = dataset
+    print(['build_faiss_index', 'end'])
 
 
 def extract_text(url: str):
@@ -25,11 +69,17 @@ def extract_text(url: str):
     return text
 
 
+def split_text(text: str):
+    lines = text.split('\n')
+    lines = [line.strip() for line in lines if line.strip()]
+    return lines
+
+
 def summarize_text(text: str):
     print(['summarize_text', 'start'])
     input_text = f'<s>Instruction: Elabora un resume del siguiente texto.\nInput: {text}\nOutput: '
     batch = tokenizer(input_text, return_tensors='pt')
-    batch = batch.to('cuda')
+    batch = batch.to(device)
     print(['summarize_text', 'generating'])
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch,
@@ -41,14 +91,49 @@ def summarize_text(text: str):
     print(['summarize_text', 'end'])
     return output
 
-def generate_question(text:str):
-    return 'Pregunta de ejemplo.'
+
+def generate_question(text: str):
+    print(['generate_question', 'start'])
+    # Get a random section of the whole text to generate a question
+    fragments = split_text(text)
+    rnd_text = random.choice(fragments)
+    shared['answer_context'] = rnd_text
+
+    input_text = f'<s>Instruction: Dado el siguiente texto quiero que generes una pregunta cuya respuesta se encuentre en él.\nInput: {rnd_text}\nOutput: '
+    batch = tokenizer(input_text, return_tensors='pt')
+    print(['generate_question', 'generating'])
+    with torch.cuda.amp.autocast():
+        output_tokens = model.generate(**batch,
+                                       max_new_tokens=256,
+                                       generation_config=generation_config)
+    output = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    print(['generate_question', 'end'])
+    return output
+
 
 def get_answer_context():
-    return 'Aquí está la respuesta.'
+    return shared['answer_context']
+
 
-def answer_question(question:str):
-    return 'Esta es la respuesta a su pregunta.'
+def answer_question(full_text: str, question: str):
+    print(['answer_question', 'start'])
+    if not shared['embeddings_dataset']:
+        build_faiss_index(full_text)
+    top_k_samples = get_nearest_examples(question, k=5)
+
+    context = '\n'.join(top_k_samples)
+
+    input_text = f"""<s>Instruction: Te voy a proporcionar un texto del cual deseo que me respondas una pregunta. 
+    El texto es el siguiente: `{context}`\nInput: {question}\nOutput: """
+    batch = tokenizer(input_text, return_tensors='pt')
+    print(['answer_question', 'generating'])
+    with torch.cuda.amp.autocast():
+        output_tokens = model.generate(**batch,
+                                       max_new_tokens=256,
+                                       generation_config=generation_config)
+    output = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    print(['answer_question', 'end'])
+    return output
 
 
 def load_model(peft_model_id):
@@ -67,5 +152,19 @@ def load_model(peft_model_id):
     return model, tokenizer
 
 
+def load_embeddings_model():
+    print(['load_embeddings_model', 'start'])
+    model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
+    print(['load_embeddings_model', 'loading tokenizer'])
+    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
+    print(['load_embeddings_model', 'loading model'])
+    model = AutoModel.from_pretrained(model_ckpt)
+    model = model.to(device)
+    print(['load_embeddings_model', 'end'])
+    return model, tokenizer
+
+
 model, tokenizer = load_model(
     "hackathon-somos-nlp-2023/opt-6.7b-lora-sag-t3000-v300-v2")
+
+emb_model, emb_tokenizer = load_embeddings_model()