Spaces:
Running
Running
File size: 2,937 Bytes
6650ee4 3c305fd 6650ee4 3c305fd 6650ee4 |
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 |
import torch
import transformers
import gradio as gr
from ragatouille import RAGPretrainedModel
from huggingface_hub import InferenceClient
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
generate_kwargs = dict(
temperature = None,
max_new_tokens = 512,
top_p = None,
do_sample = False,
)
RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
try:
gr.Info("Setting up retriever, please wait...")
_ = RAG.search("what is Mistral?", k = 1)
gr.Info("Retriever working successfully!")
except:
gr.Warning("Retriever not working!")
mark_text = '# 📚 Search Results\n'
def rag_cleaner(inp):
rank = inp['rank']
title = inp['document_metadata']['title']
content = inp['content']
return f"{rank}. <b> {title} </b> \n Abstract: {content}"
def get_prompt_text(question, context, formatted = True):
if formatted:
sys_instruction = f"Context:\n {context} \n Given the following scientific paper abstracts, take a deep breath and lets think step by step to answer what the question. Cite the titles of your sources when answering."
message = f"Question: {question}"
return f"<s>" + f"[INST] {sys_instruction} " + f" {message} [/INST] </s> "
return f"Context:\n {context} \n Given the following info, take a deep breath and lets think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"
def get_references(question, retriever, k = 10):
rag_out = retriever.search(query=question, k=k)
return rag_out
def get_rag(message):
return get_references(message, RAG)
with gr.Blocks(theme = gr.themes.Soft()) as demo:
with gr.Group():
msg = gr.Textbox(label = 'Search')
output_text = gr.Textbox(show_label = True, container = True, label = 'LLM Answer', visible = True)
input = gr.Textbox(show_label = False, visible = False)
gr_md = gr.Markdown(mark_text)
def update_with_rag_md(message):
rag_out = get_rag(message)
md_text_updated = mark_text
for i in range(10):
rag_answer = rag_out[i]
title = rag_answer['document_metadata']['title'].replace('\n','')
paper_title = f'''### [{title}](https://arxiv.org/abs/{rag_answer['document_id']})\n'''
paper_abs = rag_answer['content']
md_text_updated += paper_title + paper_abs + '\n---------------\n'+ '\n'
prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out))
return md_text_updated, prompt
def ask_llm(prompt):
output = client.text_generation(prompt, **generate_kwargs, stream=False, details=False, return_full_text=False)
output = output.lstrip(' \n') if output.lstrip().startswith('\n') else output
return gr.Textbox(output, visible = True)
msg.submit(update_with_rag_md, msg, [gr_md, input]).success(ask_llm, input, output_text)
demo.launch(debug = True)
|