Create a version for summarization using LangChain tools.

functions.py

@@ -1,11 +1,16 @@
 import os
-import requests
 import random
+
+import requests
 import torch
 from bs4 import BeautifulSoup
+from datasets import Dataset
+from langchain.docstore.document import Document
+from langchain.llms import HuggingFacePipeline
+from langchain.text_splitter import CharacterTextSplitter
 from peft import PeftConfig, PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, AutoModel
-from datasets import DatasetDict, Dataset
+from transformers import (AutoModel, AutoModelForCausalLM, AutoTokenizer,
+                          GenerationConfig, pipeline)
 
 # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
@@ -21,6 +26,8 @@ shared = {
     'full_text': None,
 }
 
+text_splitter = CharacterTextSplitter()
+
 
 def get_nearest_examples(question: str, k: int):
     print(['get_nearest_examples', 'start'])
@@ -81,7 +88,33 @@ def split_text(text: str):
     return lines
 
 
-def summarize_text(text: str):
+def remove_prompt(text: str) -> str:
+    output_prompt = 'Output: '
+    idx = text.index(output_prompt)
+    res = text[idx + len(output_prompt):].strip()
+    res = res.replace('Input: ', '')
+    return res
+
+
+def summarize_text(text: str) -> str:
+    print(['summarize_text', 'start'])
+
+    print(['summarize_text', 'splitting text'])
+    texts = text_splitter.split_text(text)
+    docs = [Document(page_content=t) for t in texts]
+    prompts = [f'<s>Instruction: Elabora un resume del siguiente texto.\nInput: {d.page_content}\nOutput: '
+               for d in docs]
+
+    print(['summarize_text', 'generating'])
+    cleaned_summaries = [remove_prompt(
+        s['generated_text']) for s in pipe(prompts)]
+    summaries = '\n\n'.join(cleaned_summaries)
+
+    print(['summarize_text', 'end'])
+    return summaries
+
+
+def summarize_text_v1(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')
@@ -145,6 +178,7 @@ def answer_question(question: str):
                                        max_new_tokens=256,
                                        generation_config=generation_config)
     output = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    output = output.replace(input_text, '')
     print(['answer_question', 'end'])
     return output
 
@@ -165,7 +199,7 @@ def load_model(peft_model_id):
     return model, tokenizer
 
 
-def load_embeddings_model(model_ckpt:str):
+def load_embeddings_model(model_ckpt: str):
     print(['load_embeddings_model', 'start'])
     print(['load_embeddings_model', 'loading tokenizer'])
     tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
@@ -176,5 +210,10 @@ def load_embeddings_model(model_ckpt:str):
     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("sentence-transformers/multi-qa-mpnet-base-dot-v1")
+model, tokenizer = load_model(
+    "hackathon-somos-nlp-2023/opt-6.7b-lora-sag-t3000-v300-v2")
+pipe = pipeline("text2text-generation", model=model,
+                tokenizer=tokenizer, max_new_tokens=100)
+llm = HuggingFacePipeline(pipeline=pipe)
+emb_model, emb_tokenizer = load_embeddings_model(
+    "sentence-transformers/multi-qa-mpnet-base-dot-v1")