wehub-resource-sync
290 lines
9.9 KiB
Python
290 lines
9.9 KiB
Python
import os
|
|
import requests
|
|
import json
|
|
import time
|
|
import io
|
|
import re
|
|
import gradio as gr
|
|
import PyPDF2
|
|
from together import Together
|
|
|
|
|
|
|
|
def download_pdf(url, save_path=None):
|
|
if url is None or 'arxiv.org' not in url:
|
|
return None
|
|
response = requests.get(url)
|
|
if save_path:
|
|
with open(save_path, 'wb') as f:
|
|
f.write(response.content)
|
|
return response.content
|
|
|
|
def extract_arxiv_pdf_url(arxiv_url):
|
|
# Check if URL is already in PDF format
|
|
if 'arxiv.org/pdf/' in arxiv_url:
|
|
return arxiv_url
|
|
|
|
# Extract arxiv_id from different URL formats
|
|
arxiv_id = None
|
|
if 'arxiv.org/abs/' in arxiv_url:
|
|
arxiv_id = arxiv_url.split('arxiv.org/abs/')[1].split()[0]
|
|
elif 'arxiv.org/html/' in arxiv_url:
|
|
arxiv_id = arxiv_url.split('arxiv.org/html/')[1].split()[0]
|
|
|
|
if arxiv_id:
|
|
return f"https://arxiv.org/pdf/{arxiv_id}.pdf"
|
|
|
|
return None # Return None if no valid arxiv_id found
|
|
|
|
def extract_text_from_pdf(pdf_content):
|
|
pdf_file = io.BytesIO(pdf_content)
|
|
reader = PyPDF2.PdfReader(pdf_file)
|
|
text = ""
|
|
for page in reader.pages:
|
|
text += page.extract_text() + "\n"
|
|
return text
|
|
|
|
def extract_references_with_llm(pdf_content):
|
|
# Extract text from PDF
|
|
text = extract_text_from_pdf(pdf_content)
|
|
|
|
# Truncate if too long
|
|
max_length = 50000
|
|
if len(text) > max_length:
|
|
text = text[:max_length] + "..."
|
|
|
|
client = Together(api_key="Your API key here")
|
|
|
|
citations = client.chat.completions.create(
|
|
model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
|
messages=[
|
|
{
|
|
"role":"user",
|
|
"content":f"Extract all the arXiv citations from Reference section of the paper including their title, authors and origins. Paper: {text} "
|
|
}
|
|
],
|
|
temperature=0.3,
|
|
)
|
|
|
|
# Prepare prompt for Llama 4
|
|
prompt = f"""
|
|
Extract the arXiv ID from the list of citations provided, including preprint arXiv ID. If there is no arXiv ID presented with the list, skip that citations.
|
|
|
|
Here are some examples on arXiv ID format:
|
|
1. arXiv preprint arXiv:1607.06450, where 1607.06450 is the arXiv ID.
|
|
2. CoRR, abs/1409.0473, where 1409.0473 is the arXiv ID.
|
|
|
|
Then, return a JSON array of objects with 'title' and 'ID' fields strictly in the following format, only return the paper title if it's arXiv ID is extracted:
|
|
|
|
Output format: [{{\"title\": \"Paper Title\", \"ID\": \"arXiv ID\"}}]
|
|
|
|
DO NOT return any other text.
|
|
|
|
List of citations:
|
|
{citations}
|
|
"""
|
|
|
|
|
|
response = client.chat.completions.create(
|
|
model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
|
messages=[
|
|
{
|
|
"role":"user",
|
|
"content":prompt
|
|
}
|
|
],
|
|
temperature=0.3,
|
|
)
|
|
response_json = response.choices[0].message.content
|
|
|
|
# Convert the JSON string to a Python object
|
|
references = []
|
|
try:
|
|
references = json.loads(response_json)
|
|
# Now you can work with `references` as a Python list or dictionary
|
|
except json.JSONDecodeError as e:
|
|
print(f"Error decoding JSON: {e}")
|
|
|
|
return references
|
|
|
|
# Check if ref_id is a valid arXiv ID
|
|
def is_valid_arxiv_id(ref_id):
|
|
# arXiv IDs are typically in the format of "1234.56789" or "1234567"
|
|
return bool(re.match(r'^\d{4}\.\d{4,5}$', ref_id) or re.match(r'^\d{7}$', ref_id))
|
|
|
|
def download_arxiv_paper_and_citations(arxiv_url, download_dir, progress=None):
|
|
if not os.path.exists(download_dir):
|
|
os.makedirs(download_dir)
|
|
|
|
if progress:
|
|
progress("Downloading main paper PDF...")
|
|
|
|
# Download main paper PDF
|
|
pdf_url = extract_arxiv_pdf_url(arxiv_url)
|
|
main_pdf_path = os.path.join(download_dir, 'main_paper.pdf')
|
|
main_pdf_content = download_pdf(pdf_url, main_pdf_path)
|
|
|
|
if main_pdf_content is None:
|
|
if progress:
|
|
progress("Invalid Url. Valid example: https://arxiv.org/abs/1706.03762v7")
|
|
return None, 0
|
|
|
|
if progress:
|
|
progress("Main paper downloaded. Extracting references...")
|
|
|
|
# Extract references using LLM
|
|
references = extract_references_with_llm(main_pdf_content)
|
|
|
|
if progress:
|
|
progress(f"Found {len(references)} references. Downloading...")
|
|
time.sleep(1)
|
|
|
|
# Download reference PDFs
|
|
all_pdf_paths = [main_pdf_path]
|
|
for i, reference in enumerate(references):
|
|
ref_title = reference.get("title")
|
|
ref_id = reference.get("ID")
|
|
if ref_id and is_valid_arxiv_id(ref_id):
|
|
ref_url = f'https://arxiv.org/pdf/{ref_id}'
|
|
ref_pdf_path = os.path.join(download_dir, f'{ref_title}.pdf')
|
|
if progress:
|
|
progress(f"Downloading reference {i+1}/{len(references)}...{ref_title}")
|
|
time.sleep(0.2)
|
|
try:
|
|
download_pdf(ref_url, ref_pdf_path)
|
|
all_pdf_paths.append(ref_pdf_path)
|
|
except Exception as e:
|
|
if progress:
|
|
progress(f"Error downloading {ref_url}: {str(e)}")
|
|
time.sleep(0.2)
|
|
|
|
# Create a list of all PDF paths
|
|
paths_file = os.path.join(download_dir, 'pdf_paths.txt')
|
|
with open(paths_file, 'w', encoding='utf-8') as f:
|
|
f.write('\n'.join(all_pdf_paths))
|
|
|
|
if progress:
|
|
progress(f"All papers downloaded. Total references: {len(references)}")
|
|
time.sleep(1)
|
|
return paths_file, len(references)
|
|
|
|
def ingest_paper_with_llama(paths_file, progress=None):
|
|
total_text = ""
|
|
total_word_count = 0
|
|
|
|
if progress:
|
|
progress("Ingesting paper content...")
|
|
|
|
with open(paths_file, 'r', encoding='utf-8') as f:
|
|
pdf_paths = f.read().splitlines()
|
|
|
|
|
|
for i, pdf_path in enumerate(pdf_paths):
|
|
if progress:
|
|
progress(f"Ingesting PDF {i+1}/{len(pdf_paths)}...")
|
|
with open(pdf_path, 'rb') as pdf_file:
|
|
pdf_content = pdf_file.read()
|
|
text = extract_text_from_pdf(pdf_content)
|
|
total_text += text + "\n\n"
|
|
total_word_count += len(text.split())
|
|
|
|
if progress:
|
|
progress("Paper ingestion complete!")
|
|
|
|
return total_text, total_word_count
|
|
|
|
def gradio_interface():
|
|
paper_content = {"text": ""}
|
|
|
|
def process(arxiv_url, progress=gr.Progress()):
|
|
download_dir = 'downloads'
|
|
progress(0, "Starting download...")
|
|
paper_path, num_references = download_arxiv_paper_and_citations(arxiv_url, download_dir,
|
|
lambda msg: progress(0.3, msg))
|
|
if paper_path is None:
|
|
return "Invalid Url. Valid example: https://arxiv.org/abs/1706.03762v7"
|
|
|
|
paper_content["text"], total_word_count = ingest_paper_with_llama(paper_path,
|
|
lambda msg: progress(0.7, msg))
|
|
progress(1.0, "Ready for chat!")
|
|
return f"Total {total_word_count} words and {num_references} reference ingested. You can now chat about the paper and citations."
|
|
|
|
def respond(message, history):
|
|
user_message = message
|
|
|
|
if not user_message:
|
|
return history, ""
|
|
|
|
# Append user message immediately
|
|
history.append([user_message, ""])
|
|
|
|
|
|
client = Together(api_key="Your API key here")
|
|
|
|
# Prepare the system prompt and user message
|
|
|
|
|
|
system_prompt = f"""
|
|
You are a research assistant that have access to the paper reference below.
|
|
Answer questions based on your knowledge on these references.
|
|
If you do not know the answer, say you don't know.
|
|
paper reference: {paper_content["text"]}
|
|
"""
|
|
|
|
stream = client.chat.completions.create(
|
|
model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_message}
|
|
],
|
|
temperature=0.3,
|
|
stream=True # Enable streaming
|
|
)
|
|
|
|
# Initialize an empty response
|
|
full_response = ""
|
|
|
|
# Stream the response chunks
|
|
for chunk in stream:
|
|
if len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None:
|
|
content = chunk.choices[0].delta.content
|
|
full_response += content
|
|
# Update the last message in history with the current response
|
|
history[-1][1] = full_response
|
|
yield history,""
|
|
|
|
|
|
|
|
def clear_chat_history():
|
|
return [], ""
|
|
|
|
with gr.Blocks(css=".orange-button {background-color: #FF7C00 !important; color: white;}") as demo:
|
|
gr.Markdown("# Research Analyzer")
|
|
with gr.Column():
|
|
input_text = gr.Textbox(label="ArXiv URL")
|
|
status_text = gr.Textbox(label="Status", interactive=False)
|
|
submit_btn = gr.Button("Ingest", elem_classes="orange-button")
|
|
submit_btn.click(fn=process, inputs=input_text, outputs=status_text)
|
|
|
|
gr.Markdown("## Chat with Llama")
|
|
chatbot = gr.Chatbot()
|
|
with gr.Row():
|
|
msg = gr.Textbox(label="Ask about the paper", scale=5)
|
|
submit_chat_btn = gr.Button("➤", elem_classes="orange-button", scale=1)
|
|
|
|
submit_chat_btn.click(respond, [msg, chatbot], [chatbot, msg])
|
|
msg.submit(respond, [msg, chatbot], [chatbot, msg])
|
|
|
|
def copy_last_response(history):
|
|
if history and len(history) > 0:
|
|
last_response = history[-1][1]
|
|
return gr.update(value=last_response)
|
|
return gr.update(value="No response to copy")
|
|
|
|
|
|
demo.launch()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
gradio_interface()
|