test_rag.py

def test_rag(pipeline, query):
  docs = vectordb.similarity_search_with_score(query)
  context = []
  for doc,score in docs:
    if(score<7):
      doc_details = doc.to_json()['kwargs']
      context.append( doc_details['page_content'])
  if(len(context)!=0):
    messages = [{"role": "user", "content": "Basándote en la siguiente información: " + "\n".join(context) + "\n Responde en castellano a la pregunta: " + query}]
    prompt = pipeline.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    outputs = pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
    answer = outputs[0]["generated_text"]
    return answer[answer.rfind("[/INST]")+8:],docs
  else:
    return "No tengo información para responder a esta pregunta",docs