Spaces:
Running
Running
| import torch | |
| import transformers | |
| import gradio as gr | |
| from ragatouille import RAGPretrainedModel | |
| from huggingface_hub import InferenceClient | |
| retrieve_results = 10 | |
| 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 = retrieve_results): | |
| 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: | |
| header = gr.Markdown("# ArXiv RAG") | |
| with gr.Group(): | |
| msg = gr.Textbox(label = 'Search') | |
| with gr.Accordion("Advanced Settings", open=False): | |
| with gr.Row(equal_height = True): | |
| llm_model = gr.Dropdown(choices = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'None'], value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model') | |
| llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results to sent as context") | |
| 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, llm_results_use = 5): | |
| rag_out = get_rag(message) | |
| md_text_updated = mark_text | |
| for i in range(retrieve_results): | |
| rag_answer = rag_out[i] | |
| title = rag_answer['document_metadata']['title'].replace('\n','') | |
| score = round(rag_answer['score'], 2) | |
| paper_title = f'''### **{score}** | [{title}](https://arxiv.org/abs/{rag_answer['document_id']})\n''' | |
| paper_abs = rag_answer['content'] | |
| authors = rag_answer['document_metadata']['authors'].replace('\n','') | |
| authors_formatted = f'*{authors}*' + ' \n\n' | |
| md_text_updated += paper_title + authors_formatted + paper_abs + '\n---------------\n'+ '\n' | |
| prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out[:llm_results_use])) | |
| return md_text_updated, prompt | |
| def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'): | |
| if llm_model_picked == 'None': | |
| return gr.Textbox(visible = False) | |
| client = InferenceClient(llm_model_picked) | |
| 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, llm_results], [gr_md, input]).success(ask_llm, [input, llm_model], output_text) | |
| demo.launch(debug = True) | |