|
|
| import os
|
| import json
|
| import concurrent.futures
|
| from mistralai.client import Mistral
|
| import arxiv
|
| from transformers import AutoTokenizer
|
| from langchain_core.prompts import PromptTemplate
|
| from langchain_text_splitters import CharacterTextSplitter
|
| from langchain_mistralai import ChatMistralAI
|
|
|
| api_key_1 = os.getenv("MISTRAL_API_KEY_1")
|
| client_1 = Mistral(api_key=api_key_1)
|
|
|
| llm = ChatMistralAI(api_key=api_key_1, model="pixtral-12b-2409")
|
|
|
| tokenizer = AutoTokenizer.from_pretrained("mistral-community/pixtral-12b")
|
|
|
| def count_tokens_in_text(text):
|
| tokens = tokenizer(text, return_tensors="pt", truncation=False, add_special_tokens=True)
|
| return len(tokens["input_ids"][0])
|
|
|
|
|
| def extract_key_topics(content, images=[]):
|
| prompt = f"""
|
| Extract the primary themes from the text below. List each theme in as few words as possible, focusing on essential concepts only. Format as a concise, unordered list with no extraneous words.
|
|
|
| ```{content}```
|
|
|
| LIST IN ENGLISH:
|
| -
|
| """
|
| message_content = [{"type": "text", "text": prompt}] + images
|
|
|
| response = client_1.chat.complete(
|
| model="pixtral-12b-2409",
|
| messages=[{"role": "user", "content": message_content}]
|
| )
|
| return response.choices[0].message.content
|
|
|
|
|
| def extract_key_topics_with_large_text(content, images=[]):
|
| text_splitter = CharacterTextSplitter.from_huggingface_tokenizer(
|
| tokenizer, chunk_size=100000, chunk_overlap=14000
|
| )
|
| docs = text_splitter.create_documents([content])
|
|
|
| image_descriptions = "\n".join(
|
| [f"Изображение {i+1}: {img['image_url']}" for i, img in enumerate(images)]
|
| )
|
|
|
| map_template = f"""
|
| Текст: {{text}}
|
| Изображения: {image_descriptions}
|
|
|
| Extract the primary themes from the text below. List each theme in as few words as possible, focusing on essential concepts only. Format as a concise, unordered list with no extraneous words.
|
| LIST IN ENGLISH:
|
| -
|
| """
|
|
|
| map_prompt = PromptTemplate.from_template(map_template)
|
| chunk_themes = []
|
|
|
| for doc in docs:
|
| formatted_prompt = map_prompt.format(text=doc.page_content)
|
| response = llm.invoke(formatted_prompt)
|
| chunk_themes.append(response.content)
|
|
|
| combined_themes = "\n".join(chunk_themes)
|
|
|
| reduce_template = f"""
|
| Следующий текст содержит несколько списков ключевых тем, извлеченных из разных частей документа:
|
| {{themes}}
|
|
|
| Extract the primary themes from the text below. List each theme in as few words as possible, focusing on essential concepts only. Format as a concise, unordered list with no extraneous words.
|
| Remove duplicates and merge similar themes.
|
|
|
| LIST IN ENGLISH:
|
| -
|
| """
|
|
|
| reduce_prompt = PromptTemplate.from_template(reduce_template)
|
| final_response = llm.invoke(reduce_prompt.format(themes=combined_themes))
|
|
|
| return final_response.content
|
|
|
|
|
| def search_relevant_articles_arxiv(key_topics, max_articles=10):
|
| articles_by_topic = {}
|
| final_topics = []
|
|
|
| def fetch_articles_for_topic(topic):
|
| topic_articles = []
|
| try:
|
| search = arxiv.Search(
|
| query=topic,
|
| max_results=max_articles,
|
| sort_by=arxiv.SortCriterion.Relevance
|
| )
|
| for result in search.results():
|
| article_data = {
|
| "title": result.title,
|
| "doi": result.doi,
|
| "summary": result.summary,
|
| "url": result.entry_id,
|
| "pdf_url": result.pdf_url
|
| }
|
| topic_articles.append(article_data)
|
| final_topics.append(topic)
|
| except Exception as e:
|
| print(f"Error fetching articles for topic '{topic}': {e}")
|
|
|
| return topic, topic_articles
|
|
|
| with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
|
| futures = {executor.submit(fetch_articles_for_topic, topic): topic for topic in key_topics}
|
| for future in concurrent.futures.as_completed(futures):
|
| topic, articles = future.result()
|
| if articles:
|
| articles_by_topic[topic] = articles
|
|
|
| return articles_by_topic, list(set(final_topics))
|
|
|
|
|
| def init(content, images=[]):
|
| if len(images) >= 8:
|
| images = images[:7]
|
|
|
| if count_tokens_in_text(content) < 128000:
|
| key_topics = extract_key_topics(content, images)
|
| else:
|
| key_topics = extract_key_topics_with_large_text(content, images)
|
|
|
| key_topics = [topic.strip("- ") for topic in key_topics.split("\n") if topic.strip()]
|
|
|
| articles_by_topic, final_topics = search_relevant_articles_arxiv(key_topics)
|
|
|
| result_json = json.dumps(articles_by_topic, indent=4)
|
|
|
| return final_topics, result_json
|
|
|
