Spaces:
Sleeping
Sleeping
| 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 | |
| logging.basicConfig(level=logging.INFO, format='[%(asctime)s][%(levelname)s] - %(message)s') | |
| # logging.getLogger().setLevel(logging.INFO) | |
| 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 | |
| embeddings = HuggingFaceEmbeddings(model_name=embeddings_model, model_kwargs={'device': 'cpu'}) | |
| if vs_hf_repo_path: | |
| hf_vectorstore = snapshot_download(repo_id=vs_hf_repo_path) | |
| self.vectore_store = FAISS.load_local(hf_vectorstore, embeddings, allow_dangerous_deserialization=True) | |
| else: | |
| self.vectore_store = FAISS.load_local(self.vectorstore_path, embeddings, 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) | |
| logging.info("Rerank model loaded!") | |
| 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) | |
| print("Inputs: ", inputs) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| logits = outputs.logits | |
| score = logits.view(-1, ).float() | |
| print("Score: ", score) | |
| return score | |
| scores = [get_score(instruction, c[0].page_content) for c in contexts] | |
| print("Scores: ", scores) | |
| 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=3): | |
| """Retrieve the most relevant contexts for a given instruction.""" | |
| logging.info("RETRIEVE DOCUMENTS") | |
| logging.info(f"Instruction: {instruction}") | |
| # Embed the query | |
| # ============================================================================================================== | |
| embedding = self.vectore_store._embed_query(instruction) | |
| logging.info(f"Query embedding generated: {len(embedding)}") | |
| # Retrieve documents | |
| # ============================================================================================================== | |
| documents_retrieved = self.vectore_store.similarity_search_with_score_by_vector(embedding, k=number_of_contexts) | |
| logging.info(f"Documents retrieved: {len(documents_retrieved)}") | |
| # Reranking | |
| # ============================================================================================================== | |
| if self.rerank_model: | |
| logging.info("RERANK DOCUMENTS") | |
| documents_reranked = self.rerank_contexts(instruction, documents_retrieved, number_of_contexts=number_of_contexts) | |
| else: | |
| logging.info("NO RERANKING") | |
| documents_reranked = documents_retrieved[:number_of_contexts] | |
| # ============================================================================================================== | |
| return documents_reranked | |
| 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"]) | |
| response = "" | |
| for i, (doc, score) in enumerate(docs): | |
| response += "\n\n" + "="*100 | |
| response += f"\nDocument {i+1}" | |
| response += "\n" + "="*100 | |
| response += f"\nScore: {score:.5f}" | |
| response += f"\nTitle: {doc.metadata['title']}" | |
| response += f"\nURL: {doc.metadata['url']}" | |
| response += f"\nID: {doc.metadata['id']}" | |
| response += f"\nStart index: {doc.metadata['start_index']}" | |
| # response += f"\nSource: {doc.metadata['src']}" | |
| # response += f"\nRedirected: {doc.metadata['redirected']}" | |
| # url = doc.metadata['url'] | |
| # response += f"\nRevision ID: {url}" | |
| # response += f'\nURL: <a href="{url}" target="_blank">{url}</a><br>' | |
| response += "\n" + "-"*100 + "\n" | |
| response += f"\nContent:\n" | |
| response += doc.page_content | |
| full_context = "" | |
| source = [] | |
| if not response: | |
| return self.NO_ANSWER_MESSAGE | |
| return response, full_context, source | |
| except Exception as err: | |
| print(err) | |