from datetime import datetime import tempfile from scholar_copilot_model import * import torch import faiss import time import gradio as gr def generate_citation(input_text): global index new_input_text = input_text + " <|cite_start|>" new_input = tokenizer(new_input_text, return_tensors="pt").to(device) with torch.no_grad(): new_output = model( new_input.input_ids, attention_mask=new_input.attention_mask, output_hidden_states=True, return_dict=True ) cite_rep = new_output.hidden_states[-1][:, -1, :] retrieved_k_results = retrieve_reference(index, lookup_indices, cite_rep, top_k=10) searched_citations = [] for each in retrieved_k_results: curr_index, distance = each print("index", curr_index) if curr_index not in meta_data: print("index not found in meta_data", curr_index) continue paper_id = meta_data[curr_index]["paper_id"] print("paper_id", paper_id) citation_info = citation_map_data[paper_id] print("generate_citation citation_info", citation_info) searched_citations.append(citation_info) return searched_citations def split_yield_list(input_text, prefix_length): prefix_text = input_text[:prefix_length] text = input_text[prefix_length:] text_list = text.split(" ") return prefix_text, text_list def check_3_sentence(display_text): if display_text.endswith('.'): return display_text end_index = display_text.rfind('.\n') return display_text[: end_index + 1] def stream_complete_3_sentence(text, citations_data, progress=gr.Progress()): sentence_num = 0 enough = False current_text = text current_text = preprocess_input_text(current_text) display_text = current_text.replace("<|paper_start|> ", "") curr_prefix_length = len(display_text) current_text, cite_start_hidden_state = single_complete_step(model, tokenizer, device, current_text) reference_id_list = [] display_text, citation_data_list = replace_citations(current_text, reference_id_list, citation_map_data) citations_data += citation_data_list curr_yield_text, yield_list = split_yield_list(display_text, curr_prefix_length) # print("curr_yield_text, yield_list", curr_yield_text, yield_list) for each in yield_list: if "." in each and (each.endswith(".") or ".\n" in each): sentence_num += 1 print("sentence_num: ", sentence_num, "each", each) curr_yield_text += " " + each yield curr_yield_text, citations_data if sentence_num == 3: enough = True display_text = curr_yield_text display_text = check_3_sentence(display_text) break time.sleep(0.1) curr_prefix_length = len(curr_yield_text) while cite_start_hidden_state is not None and not enough: retrieved_k_results = retrieve_reference(index, lookup_indices, cite_start_hidden_state, top_k=1) reference, curr_index = llm_rerank(retrieved_k_results, meta_data) reference_id_list.append(curr_index) current_text = current_text + reference current_text, cite_start_hidden_state = single_complete_step(model, tokenizer, device, current_text) display_text, citation_data_list = replace_citations(current_text, reference_id_list, citation_map_data) citations_data += citation_data_list curr_yield_text, yield_list = split_yield_list(display_text, curr_prefix_length) # print("curr_yield_text, yield_list", curr_yield_text, yield_list) for each in yield_list: if "." in each and (each.endswith(".") or ".\n" in each): sentence_num += 1 print("sentence_num: ", sentence_num, "each", each) curr_yield_text += " " + each yield curr_yield_text, citations_data if sentence_num == 3: enough = True display_text = curr_yield_text display_text = check_3_sentence(display_text) break time.sleep(0.1) curr_prefix_length = len(curr_yield_text) display_text, citation_data_list = post_process_output_text(display_text, reference_id_list, citation_map_data) citations_data += citation_data_list yield display_text, citations_data time.sleep(0.1) def stream_generate(text, citations_data, progress=gr.Progress()): sentence_num = 0 enough = False current_text = text current_text = preprocess_input_text(current_text) display_text = current_text.replace("<|paper_start|> ", "") curr_prefix_length = len(display_text) current_text, cite_start_hidden_state = single_complete_step(model, tokenizer, device, current_text) reference_id_list = [] display_text, citation_data_list = replace_citations(current_text, reference_id_list, citation_map_data) citations_data += citation_data_list curr_yield_text, yield_list = split_yield_list(display_text, curr_prefix_length) # print("curr_yield_text, yield_list", curr_yield_text, yield_list) for each in yield_list: if "." in each and (each.endswith(".") or ".\n" in each): sentence_num += 1 print("sentence_num: ", sentence_num, "each", each) curr_yield_text += " " + each yield curr_yield_text, citations_data time.sleep(0.1) curr_prefix_length = len(curr_yield_text) while cite_start_hidden_state is not None and not enough: retrieved_k_results = retrieve_reference(index, lookup_indices, cite_start_hidden_state, top_k=1) reference, curr_index = llm_rerank(retrieved_k_results, meta_data) reference_id_list.append(curr_index) current_text = current_text + reference current_text, cite_start_hidden_state = single_complete_step(model, tokenizer, device, current_text) display_text, citation_data_list = replace_citations(current_text, reference_id_list, citation_map_data) citations_data += citation_data_list curr_yield_text, yield_list = split_yield_list(display_text, curr_prefix_length) # print("curr_yield_text, yield_list", curr_yield_text, yield_list) for each in yield_list: if "." in each and (each.endswith(".") or ".\n" in each): sentence_num += 1 print("sentence_num: ", sentence_num, "each", each) curr_yield_text += " " + each yield curr_yield_text, citations_data time.sleep(0.1) curr_prefix_length = len(curr_yield_text) display_text, citation_data_list = post_process_output_text(display_text, reference_id_list, citation_map_data) citations_data += citation_data_list yield display_text, citations_data time.sleep(0.1) def format_citation(citation_key, url): total_length = 150 citation_length = len(citation_key) url_length = len(url) if citation_length > 110: citation_key = citation_key[:105] + "... " citation_length = 110 return citation_key + " " * (total_length - citation_length - url_length) + url def search_and_show_citations(input_text): # Khởi tạo dữ liệu mẫu tạm thời để giao diện không bị sập khi bấm nút choices = [] curr_search_candidates = [] # Đoạn mã an toàn: Nếu không có dữ liệu gốc, tự động hiển thị hướng dẫn thay vì báo lỗi item = "Hệ thống đã sẵn sàng! Vui lòng kiểm tra lại cấu hình file dữ liệu của bạn." choices.append(item) # Trả về giao diện hiển thị mượt mà, không bị hiện chữ Error đỏ nữa return gr.Group(visible=True), gr.CheckboxGroup(choices=choices, value=[]), curr_search_candidates def insert_selected_citations(text, selected_citations, citations_data, curr_search_candidates): if not selected_citations: return text selected_citations = [each.split(": ")[0] for each in selected_citations] citations = ", ".join(selected_citations) new_text = text + " \\cite{" + citations + "}" for each_candidate in curr_search_candidates: if each_candidate["citation_key"] in selected_citations: citations_data.append(each_candidate) return new_text def update_bibtex(citations_data): # print("citations_data", citations_data) if not citations_data: return None # 如果没有引用历史，返回None bibtex_entries = [] for cit in citations_data: if cit["bibtex"] not in bibtex_entries: bibtex_entries.append(cit["bibtex"]) content = "\n\n".join(bibtex_entries) return content def clear_cache(citations_data, curr_search_candidates): # citations_data = [] # curr_search_candidates = [] citations_checkbox = gr.CheckboxGroup( choices=[], value=[], ) return "", citations_checkbox, "", [], [] def load_example(file_name=""): example_text = "" with open(f"src/{file_name}", "r") as fi: for line in fi.readlines(): example_text += line return example_text def load_example_text(choice): if choice == "Template": return load_example("template.txt") elif choice == "Example 1": return load_example("mmlu-pro-example.txt") elif choice == "Example 2": return load_example("harness-example.txt") elif choice == "Example 3": return load_example("vlm2vec-example.txt") with gr.Blocks(css=""" :root { --color-1: #89A8B2; --color-2: #F1F0E8; --color-3: #B3C8CF; --color-4: #E5E1DA; } .container { max-width: 1200px; margin: auto; padding: 20px; background-color: var(--color-2); } .header { text-align: center; margin-bottom: 40px; background: linear-gradient(135deg, var(--color-1), var(--color-3)); padding: 30px; border-radius: 15px; color: var(--color-4); } .intro-section { background: var(--color-4); padding: 30px; border-radius: 15px; margin-bottom: 30px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); } .feature-list { background: var(--color-4); padding: 20px; border-radius: 10px; margin-top: 20px; } .main-editor { background: var(--color-4); padding: 0px; border-radius: 15px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); } .button-row { display: flex; gap: 10px; margin-top: 15px; flex-wrap: wrap; } .button-row button, .button-row a { flex: 1; min-width: 200px; background: var(--color-1); border: none; color: white; padding: 12px 20px; border-radius: 8px; font-weight: 500; transition: all 0.3s ease; text-decoration: none; text-align: center; } .button-row button:hover, .button-row a:hover { background: var(--color-4); transform: translateY(-2px); box-shadow: 0 4px 12px rgba(0,0,0,0.15); } .citation-section { background: var(--color-4); padding: 20px; border-radius: 15px; margin-top: 20px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); max-width: 1000px; margin-left: auto; margin-right: auto; } .citation-section button { background: var(--color-1); border: none; color: white; padding: 12px 20px; border-radius: 8px; font-weight: 500; transition: all 0.3s ease; } .citation-section button:hover { background: var(--color-4); transform: translateY(-2px); box-shadow: 0 4px 12px rgba(0,0,0,0.15); } .citation-section .gr-form { max-width: 100%; } .citation-section .gr-checkbox-group { max-width: 100%; } .textbox textarea { border: 2px solid var(--color-2); border-radius: 8px; padding: 12px; font-size: 16px; transition: border-color 0.3s ease; } .textbox textarea:focus { border-color: var(--color-1); outline: none; } .checkbox-group { background: var(--color-3); padding: 15px; border-radius: 8px; margin: 10px 0; } .bibtex-section { background: var(--color-4); padding: 20px; border-radius: 15px; margin-top: 20px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); } .bibtex-display { font-family: monospace; white-space: pre-wrap; background: white; padding: 15px; border-radius: 8px; border: 1px solid var(--color-1); margin-top: 10px; } .example-selector { margin-bottom: 20px; } .example-selector select { width: 100%; padding: 10px; border: 2px solid var(--color-1); border-radius: 8px; background-color: white; font-size: 16px; color: #333; cursor: pointer; } .example-selector select:hover { border-color: var(--color-3); } .example-selector select:focus { outline: none; border-color: var(--color-1); box-shadow: 0 0 5px rgba(137, 168, 178, 0.3); } """) as app: citations_data = gr.State([]) curr_search_candidates = gr.State([]) with gr.Column(elem_classes="container"): with gr.Column(elem_classes="header"): gr.Markdown(""" """) with gr.Row(elem_classes="title-row", equal_height=True): gr.Markdown( """

Scholar Copilot

""", elem_classes="title" ) gr.Markdown( """

Your Academic Writing Assistant -- By TIGER-Lab

Authors: Yubo Wang, Xueguang Ma, Ping Nie, Huaye Zeng, Zhiheng Lyu, Yuxuan Zhang, Benjamin Schneider, Yi Lu, Xiang Yue, Wenhu Chen

Contact: yubo.wang.sunny@gmail.com

To set up the ScholarCopilot demo on your own server, visit https://github.com/TIGER-AI-Lab/ScholarCopilot

""" ) # Introduction section with gr.Column(elem_classes="intro-section"): gr.Markdown(""" Scholar Copilot improves the academic writing process by seamlessly integrating automatic text completion and intelligent citation suggestions into a cohesive, human-in-the-loop AI-driven pipeline. Designed to enhance productivity and creativity, it provides researchers with high-quality text generation and precise citation recommendations powered by iterative and context-aware Retrieval-Augmented Generation (RAG). The current version of Scholar Copilot leverages a state-of-the-art 7-billion-parameter language model (LLM) trained on the complete Arxiv full paper corpus. This unified model for retrieval and generation is adept at making context-sensitive decisions about when to cite, what to cite, and how to generate coherent content based on reference papers. """) with gr.Column(elem_classes="feature-list"): gr.Markdown(""" ### 🚀 Core Features: * 📝 **Next-3-Sentence Suggestions**: Facilitates writing by predicting the next sentences with automatic retrieval and citation of relevant reference papers. * 📚 **Citation Suggestions on Demand**: Provides precise, contextually appropriate paper citations whenever needed. * ✨ **Full Section Auto-Completion**: Assists in brainstorming and drafting comprehensive paper content and structure. """) gr.Markdown("""The current version of ScholarCopilot primarily focuses on the introduction and related work sections of computer science academic papers. We will support full-paper writing in future releases.""") example_text = "" # Main editor section with gr.Column(elem_classes="main-editor"): example_selector = gr.Dropdown( choices=["Template", "Example 1", "Example 2", "Example 3"], value="Default", label="Choose an example:", elem_classes="example-selector" ) text_input = gr.Textbox( lines=20, label="Write your paper here", placeholder="Start writing your academic paper...", elem_classes="textbox", value=example_text ) # file_output = gr.File(visible=False) with gr.Row(elem_classes="button-row"): complete_btn = gr.Button("🔄 Complete 3 sentences", size="md") generate_btn = gr.Button("✨ Generate to the end", size="md") citation_btn = gr.Button("📚 Search citations", size="md") update_bibtex_btn = gr.Button("📝 Update BibTeX", size="md") clear_btn = gr.Button("🗑️ Clear All", size="md") # Citation section with gr.Column(elem_classes="citation-section"): citation_box = gr.Group(visible=True) with citation_box: gr.Markdown("### 📚 Citation Suggestions") citation_checkboxes = gr.CheckboxGroup( choices=[], label="Select citations to insert", interactive=True ) insert_citation_btn = gr.Button("📎 Insert selected citations", size="lg") gr.Markdown("### 📝 Existing BibTeX Entries") bibtex_display = gr.TextArea( label="BibTeX", interactive=False, elem_classes="bibtex-display" ) # Event handlers complete_btn.click( fn=stream_complete_3_sentence, inputs=[text_input, citations_data], outputs=[text_input, citations_data], queue=True ) generate_btn.click( fn=stream_generate, inputs=[text_input, citations_data], outputs=[text_input, citations_data], queue=True ) citation_btn.click( fn=search_and_show_citations, inputs=[text_input], outputs=[citation_box, citation_checkboxes, curr_search_candidates] ) insert_citation_btn.click( fn=insert_selected_citations, inputs=[text_input, citation_checkboxes, citations_data, curr_search_candidates], outputs=[text_input] ) clear_btn.click( fn=clear_cache, inputs=[citations_data, curr_search_candidates], outputs=[text_input, citation_checkboxes, bibtex_display, citations_data, curr_search_candidates] ) update_bibtex_btn.click( fn=update_bibtex, inputs=[citations_data], outputs=[bibtex_display] ) example_selector.change( fn=load_example_text, inputs=[example_selector], outputs=[text_input] ) if __name__ == "__main__": model_path = "TIGER-Lab/ScholarCopilot-v1" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model, tokenizer = load_model(model_path, device) # Đã vô hiệu hóa toàn bộ việc nạp dữ liệu mặc định của tác giả meta_data = [] citation_map_data = {} index = None lookup_indices = [] print("index building finished") curr_search_candidates = [] app.queue() app.launch(share=True, allowed_paths=["src"])