|
from functools import lru_cache |
|
import time, aiohttp, asyncio, json, os, multiprocessing, torch, \ |
|
requests, xmltodict, fitz, io |
|
from minivectordb.embedding_model import EmbeddingModel |
|
from minivectordb.vector_database import VectorDatabase |
|
from text_util_en_pt.cleaner import structurize_text, detect_language, Language |
|
import gradio as gr |
|
|
|
torch.set_num_threads(2) |
|
|
|
openrouter_key = os.environ.get("OPENROUTER_KEY") |
|
model = EmbeddingModel(use_quantized_onnx_model=True) |
|
|
|
def convert_xml_to_json(xml): |
|
return xmltodict.parse(xml) |
|
|
|
def clean_title(title): |
|
title = title.replace('\n', ' ') |
|
while ' ' in title: |
|
title = title.replace(' ', ' ') |
|
return title |
|
|
|
@lru_cache(maxsize=500) |
|
def fetch_arxiv_links(query, max_results=5): |
|
url = f'http://export.arxiv.org/api/query?search_query=all:{query}&start=0&max_results={max_results}' |
|
response = requests.get(url) |
|
json_response = convert_xml_to_json(response.text) |
|
|
|
|
|
entries = [] |
|
for entry in json_response['feed']['entry']: |
|
|
|
title = entry['title'] |
|
id = entry['id'].split('/abs/')[-1] |
|
|
|
link = f'http://arxiv.org/abs/{id}' |
|
pdf_link = f'http://arxiv.org/pdf/{id}.pdf' |
|
|
|
entries.append({ |
|
'title': clean_title(title), |
|
'link': link, |
|
'pdf_link': pdf_link |
|
}) |
|
return entries |
|
|
|
def download_pdf_from_link(link): |
|
|
|
response = requests.get(link) |
|
return io.BytesIO(response.content) |
|
|
|
@lru_cache(maxsize=100) |
|
def read_remote_pdf(pdf_metadata): |
|
pdf_metadata = json.loads(pdf_metadata) |
|
|
|
link = pdf_metadata['pdf_link'] |
|
title = pdf_metadata['title'] |
|
|
|
pdf_content = download_pdf_from_link(link) |
|
pdf_file = fitz.open("pdf", pdf_content.read()) |
|
text_content = [page.get_text() for page in pdf_file] |
|
pdf_file.close() |
|
del pdf_file |
|
return {'title': title, 'text': '\n'.join(text_content)} |
|
|
|
def fetch_data_from_pdfs(links): |
|
links = [ json.dumps(link) for link in links ] |
|
with multiprocessing.Pool(10) as pool: |
|
pdf_metadata = pool.map(read_remote_pdf, links) |
|
return pdf_metadata |
|
|
|
def index_and_search(query, pdf_metadata): |
|
start = time.time() |
|
query_embedding = model.extract_embeddings(query) |
|
|
|
|
|
vector_db = VectorDatabase() |
|
|
|
sentence_counter = 1 |
|
|
|
for pdf_data in pdf_metadata: |
|
text = pdf_data['text'] |
|
title = pdf_data['title'] |
|
|
|
sentences = [ s['sentence'] for s in structurize_text(text)] |
|
|
|
for sentence in sentences: |
|
sentence_embedding = model.extract_embeddings(sentence) |
|
vector_db.store_embedding( |
|
sentence_counter, |
|
sentence_embedding, |
|
{ |
|
'sentence': sentence, |
|
'title': title |
|
} |
|
) |
|
sentence_counter += 1 |
|
|
|
embedding_time = time.time() - start |
|
|
|
|
|
start = time.time() |
|
search_results = vector_db.find_most_similar(query_embedding, k = 15) |
|
search_metadata = search_results[2] |
|
retrieval_time = time.time() - start |
|
|
|
retrieved_contents = {} |
|
for ret_cont in search_metadata: |
|
title = ret_cont['title'] |
|
if title not in retrieved_contents: |
|
retrieved_contents[title] = [] |
|
retrieved_contents[title].append(ret_cont['sentence']) |
|
|
|
retrieved_contents = {k: '\n'.join(v) for k, v in retrieved_contents.items() if len(v) > 2} |
|
|
|
return retrieved_contents, embedding_time, retrieval_time |
|
|
|
def retrieval_pipeline(query, question): |
|
start = time.time() |
|
links = fetch_arxiv_links(query) |
|
websearch_time = time.time() - start |
|
|
|
start = time.time() |
|
pdf_metadata = fetch_data_from_pdfs(links) |
|
webcrawl_time = time.time() - start |
|
|
|
retrieved_contents, embedding_time, retrieval_time = index_and_search(question, pdf_metadata) |
|
|
|
return retrieved_contents, websearch_time, webcrawl_time, embedding_time, retrieval_time, links |
|
|
|
|
|
async def predict(message, history): |
|
|
|
|
|
message = message.split(';') |
|
|
|
query = message[0].split(':')[-1].strip() |
|
question = message[1].split(':')[-1].strip() |
|
|
|
retrieved_contents, websearch_time, webcrawl_time, embedding_time, retrieval_time, links = retrieval_pipeline(query, question) |
|
|
|
if detect_language(message) == Language.ptbr: |
|
context = "" |
|
for title, content in retrieved_contents.items(): |
|
context += f'Artigo "{title}"\nConteúdo:\n{content}\n\n' |
|
prompt = f'{context.strip()}\n\nBaseado nos conteúdos dos artigos, responda: "{question}"\n\nPor favor, mencione a fonte da sua resposta.\nResponda somente em português brasileiro' |
|
else: |
|
context = "" |
|
for title, content in retrieved_contents.items(): |
|
context += f'Article "{title}"\nContent:\n{content}\n\n' |
|
prompt = f'{context.strip()}\n\nBased on the article\'s contents, answer: "{question}"\n\nPlease, mention the source of your answer.' |
|
|
|
print(prompt) |
|
|
|
url = "https://openrouter.ai/api/v1/chat/completions" |
|
headers = { "Content-Type": "application/json", |
|
"Authorization": f"Bearer {openrouter_key}" } |
|
body = { "stream": True, |
|
"model": "deepseek/deepseek-chat", |
|
"max_tokens": 1024, |
|
"messages": [ |
|
{"role": "user", "content": prompt} |
|
] } |
|
|
|
full_response = "" |
|
async with aiohttp.ClientSession() as session: |
|
async with session.post(url, headers=headers, json=body) as response: |
|
buffer = "" |
|
async for chunk in response.content.iter_any(): |
|
buffer += chunk.decode() |
|
while "\n" in buffer: |
|
line, buffer = buffer.split("\n", 1) |
|
|
|
if line.startswith("data: "): |
|
event_data = line[len("data: "):] |
|
if event_data != '[DONE]': |
|
try: |
|
current_text = json.loads(event_data)['choices'][0]['delta']['content'] |
|
full_response += current_text |
|
yield full_response |
|
await asyncio.sleep(0.01) |
|
except Exception: |
|
try: |
|
current_text = json.loads(event_data)['choices'][0]['text'] |
|
full_response += current_text |
|
yield full_response |
|
await asyncio.sleep(0.01) |
|
except Exception: |
|
pass |
|
|
|
final_metadata_block = "" |
|
|
|
final_metadata_block += f"Links visited:\n" |
|
for link in links: |
|
final_metadata_block += f"{link['title']} ({link['link']})\n" |
|
final_metadata_block += f"\nWeb search time: {websearch_time:.4f} seconds\n" |
|
final_metadata_block += f"\nText extraction: {webcrawl_time:.4f} seconds\n" |
|
final_metadata_block += f"\nEmbedding time: {embedding_time:.4f} seconds\n" |
|
final_metadata_block += f"\nRetrieval from VectorDB time: {retrieval_time:.4f} seconds" |
|
|
|
yield f"{full_response}\n\n{final_metadata_block}" |
|
|
|
gr.ChatInterface( |
|
predict, |
|
title="Automated Arxiv Paper Search and Question Answering", |
|
description="Provide a search term and a question to find relevant papers and answer questions about them.", |
|
retry_btn=None, |
|
undo_btn=None, |
|
examples=[ |
|
'Search: RAG LLM; Question: What are some challenges of implementing a system of RAG with LLMs?', |
|
'Search: LLM Self-Play; Question: What are the benefits of using self-play with LLMs?', |
|
'Search: Portable Blockchain; Question: How can a portable blockchain device be implemented?', |
|
'Search: 1.58 bit LLMs; Question: How do 1.58 bit LLMs work? Is there an available model to test?', |
|
'Search: Programação Robocode; Question: Como posso utilizar o robocode no contexto de aprendizagem de programação?', |
|
'Search: Pensamento Computacional; Question: Explique os conceitos do pensamento computacional.' |
|
] |
|
).launch() |
