Spaces:
Runtime error
Runtime error
File size: 5,507 Bytes
b163a23 f6691f3 b163a23 f0a0564 b163a23 |
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 |
import logging
import os
import requests
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from openai import OpenAI
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 = " renfevectors"
def __init__(self, hf_token, embeddings_model, model_name, rerank_model, rerank_number_contexts):
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
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model, model_kwargs={'device': 'cpu'})
self.vectore_store = FAISS.load_local(self.vectorstore, embeddings, allow_dangerous_deserialization=True)#, allow_dangerous_deserialization=True)
logging.info("RAG loaded!")
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."""
documentos = self.vectore_store.similarity_search_with_score(instruction, k=self.rerank_number_contexts)
documentos = self.rerank_contexts(instruction, documentos, number_of_contexts=number_of_contexts)
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": instruction}
],
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:
text_context += doc[0].page_content
full_context += doc[0].page_content + "\n"
full_context += doc[0].metadata["Fuente"] + "\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)
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)
|