Marina Pliusnina
change default parameters
f5848c0
raw
history blame
1.9 kB
import logging
import os
import requests
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."
def __init__(self, hf_token, embeddings_model, model_name):
self.model_name = model_name
self.hf_token = hf_token
# load vectore store
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model, model_kwargs={'device': 'cpu'})
self.vectore_store = FAISS.load_local("vectorestore", embeddings, allow_dangerous_deserialization=True)#, allow_dangerous_deserialization=True)
logging.info("RAG loaded!")
def get_context(self, instruction, number_of_contexts=3):
context = ""
documentos = self.vectore_store.similarity_search_with_score(instruction, k=number_of_contexts)
for doc in documentos:
context += doc[0].page_content
return context
def predict(self, instruction, context):
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 "
payload = {
"inputs": query,
"parameters": {"MAX_NEW_TOKENS": 750, "TEMPERATURE": 0.35}
}
response = requests.post(self.model_name, headers=headers, json=payload)
return response.json()[0]["generated_text"].split("###")[-1][8:-1]
def get_response(self, prompt: str) -> str:
context = self.get_context(prompt)
response = self.predict(prompt, context)
if not response:
return self.NO_ANSWER_MESSAGE
return response