|
from transformers import pipeline |
|
|
|
repo_id = "jonruida/model-IC" |
|
|
|
query_pipeline = transformers.pipeline( |
|
"text-generation", |
|
model=model, |
|
tokenizer=tokenizer, |
|
torch_dtype=torch.float16, |
|
device_map="auto", max_new_tokens=200) |
|
|
|
def test_rag(pipeline, input_text): |
|
docs = chroma_db/chroma.sqlite3.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 |
|
|
|
|
