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") mark_text = '# 📚 Search Results\n' def rag_cleaner(inp): rank = inp['rank'] title = inp['document_metadata']['title'] content = inp['content'] return f"{rank}. {title} \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"" + f"[INST] {sys_instruction} " + f" {message} [/INST] " 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)