Spaces:
Sleeping
Sleeping
File size: 6,507 Bytes
f75ccae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
import logging
import os
import requests
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from openai import OpenAI
from huggingface_hub import snapshot_download, InferenceClient
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
class RAG:
NO_ANSWER_MESSAGE: str = "Ho sento, no he pogut respondre la teva pregunta."
#vectorstore = "index-intfloat_multilingual-e5-small-500-100-CA-ES" # mixed
#vectorstore = "vectorestore" # CA only
#vectorstore = "index-BAAI_bge-m3-1500-200-recursive_splitter-CA_ES_UE"
def __init__(self, vs_hf_repo_path, vectorstore_path, hf_token, embeddings_model, model_name, rerank_model, rerank_number_contexts):
self.vs_hf_repo_path = vs_hf_repo_path
self.vectorstore_path=vectorstore_path
self.model_name = model_name
self.hf_token = hf_token
self.rerank_model = rerank_model
self.rerank_number_contexts = rerank_number_contexts
# load vectore store
hf_vectorstore = snapshot_download(repo_id=vs_hf_repo_path)
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model, model_kwargs={'device': 'cpu'})
self.vectore_store = FAISS.load_local(hf_vectorstore, embeddings, allow_dangerous_deserialization=True)
# self.vectore_store = FAISS.load_local(self.vectorstore_path, embeddings, allow_dangerous_deserialization=True)#, allow_dangerous_deserialization=True)
logging.info("RAG loaded!")
logging.info( self.vectore_store)
def rerank_contexts(self, instruction, contexts, number_of_contexts=1):
"""
Rerank the contexts based on their relevance to the given instruction.
"""
rerank_model = self.rerank_model
tokenizer = AutoTokenizer.from_pretrained(rerank_model)
model = AutoModelForSequenceClassification.from_pretrained(rerank_model)
def get_score(query, passage):
"""Calculate the relevance score of a passage with respect to a query."""
inputs = tokenizer(query, passage, return_tensors='pt', truncation=True, padding=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
score = logits.view(-1, ).float()
return score
scores = [get_score(instruction, c[0].page_content) for c in contexts]
combined = list(zip(contexts, scores))
sorted_combined = sorted(combined, key=lambda x: x[1], reverse=True)
sorted_texts, _ = zip(*sorted_combined)
return sorted_texts[:number_of_contexts]
def get_context(self, instruction, number_of_contexts=2):
"""Retrieve the most relevant contexts for a given instruction."""
logging.info("RETRIEVE DOCUMENTS")
documentos = self.vectore_store.similarity_search_with_score(instruction, k=self.rerank_number_contexts)
# logging.info(documentos)
logging.info("RERANK DOCUMENTS")
documentos = self.rerank_contexts(instruction, documentos, number_of_contexts=number_of_contexts)
# logging.info(documentos)
print("Reranked documents")
return documentos
def predict_dolly(self, instruction, context, model_parameters):
api_key = os.getenv("HF_TOKEN")
headers = {
"Accept" : "application/json",
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
query = f"### Instruction\n{instruction}\n\n### Context\n{context}\n\n### Answer\n "
#prompt = "You are a helpful assistant. Answer the question using only the context you are provided with. If it is not possible to do it with the context, just say 'I can't answer'. <|endoftext|>"
payload = {
"inputs": query,
"parameters": model_parameters
}
response = requests.post(self.model_name, headers=headers, json=payload)
return response.json()[0]["generated_text"].split("###")[-1][8:]
def predict_completion(self, instruction, context, model_parameters):
client = OpenAI(
base_url=os.getenv("MODEL"),
api_key=os.getenv("HF_TOKEN")
)
query = f"Context:\n{context}\n\nQuestion:\n{instruction}"
chat_completion = client.chat.completions.create(
model="tgi",
messages=[
{"role": "user", "content": query}
],
temperature=model_parameters["temperature"],
max_tokens=model_parameters["max_new_tokens"],
stream=False,
stop=["<|im_end|>"],
extra_body = {
"presence_penalty": model_parameters["repetition_penalty"] - 2,
"do_sample": False
}
)
response = chat_completion.choices[0].message.content
return response
def beautiful_context(self, docs):
text_context = ""
full_context = ""
source_context = []
for doc in docs:
# print("="*100)
# logging.info(doc)
text_context += doc[0].page_content
full_context += doc[0].page_content + "\n"
full_context += doc[0].metadata["title"] + "\n\n"
full_context += doc[0].metadata["url"] + "\n\n"
source_context.append(doc[0].metadata["url"])
return text_context, full_context, source_context
def get_response(self, prompt: str, model_parameters: dict) -> str:
try:
docs = self.get_context(prompt, model_parameters["NUM_CHUNKS"])
text_context, full_context, source = self.beautiful_context(docs)
print("#"*100)
logging.info("text_context")
logging.info(text_context)
print("#"*100)
logging.info("full context")
logging.info(full_context)
print("#"*100)
logging.info("source")
logging.info(source)
del model_parameters["NUM_CHUNKS"]
response = self.predict_completion(prompt, text_context, model_parameters)
if not response:
return self.NO_ANSWER_MESSAGE
return response, full_context, source
except Exception as err:
print(err)
|