import streamlit as st from groq import Groq import json import os from io import BytesIO from markdown import markdown from weasyprint import HTML, CSS from dotenv import load_dotenv import fitz # PyMuPDF import faiss from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity import numpy as np import tempfile # Load .env file to environment load_dotenv() GROQ_API_KEY = os.getenv("GROQ_API_KEY", None) if "api_key" not in st.session_state: st.session_state.api_key = GROQ_API_KEY if "groq" not in st.session_state: if GROQ_API_KEY: st.session_state.groq = Groq() class GenerationStatistics: def __init__( self, input_time=0, output_time=0, input_tokens=0, output_tokens=0, total_time=0, model_name="mixtral-8x7b-32768", ): self.input_time = input_time self.output_time = output_time self.input_tokens = input_tokens self.output_tokens = output_tokens self.total_time = ( total_time # Sum of queue, prompt (input), and completion (output) times ) self.model_name = model_name def get_input_speed(self): """ Tokens per second calculation for input """ if self.input_time != 0: return self.input_tokens / self.input_time else: return 0 def get_output_speed(self): """ Tokens per second calculation for output """ if self.output_time != 0: return self.output_tokens / self.output_time else: return 0 def add(self, other): """ Add statistics from another GenerationStatistics object to this one. """ if not isinstance(other, GenerationStatistics): raise TypeError("Can only add GenerationStatistics objects") self.input_time += other.input_time self.output_time += other.output_time self.input_tokens += other.input_tokens self.output_tokens += other.output_tokens self.total_time += other.total_time def __str__(self): return ( f"\n## {self.get_output_speed():.2f} T/s ⚡\nRound trip time: {self.total_time:.2f}s Model: {self.model_name}\n\n" f"| Metric | Input | Output | Total |\n" f"|-----------------|----------------|-----------------|----------------|\n" f"| Speed (T/s) | {self.get_input_speed():.2f} | {self.get_output_speed():.2f} | {(self.input_tokens + self.output_tokens) / self.total_time if self.total_time != 0 else 0:.2f} |\n" f"| Tokens | {self.input_tokens} | {self.output_tokens} | {self.input_tokens + self.output_tokens} |\n" f"| Inference Time (s) | {self.input_time:.2f} | {self.output_time:.2f} | {self.total_time:.2f} |" ) class ResearchPaper: def __init__(self, paper_title, structure): self.paper_title = paper_title self.structure = structure self.contents = {title: "" for title in self.flatten_structure(structure)} self.placeholders = {title: st.empty() for title in self.flatten_structure(structure)} st.markdown(f"# {self.paper_title}") st.markdown("## Generating the following:") toc_columns = st.columns(4) self.display_toc(self.structure, toc_columns) st.markdown("---") def flatten_structure(self, structure): sections = [] for title, content in structure.items(): sections.append(title) if isinstance(content, dict): sections.extend(self.flatten_structure(content)) return sections def update_content(self, title, new_content): try: self.contents[title] += new_content self.display_content(title) except TypeError as e: pass def display_content(self, title): if self.contents[title].strip(): self.placeholders[title].markdown(f"## {title}\n{self.contents[title]}") def display_structure(self, structure=None, level=1): if structure is None: structure = self.structure for title, content in structure.items(): if self.contents[title].strip(): # Only display title if there is content st.markdown(f"{'#' * level} {title}") self.placeholders[title].markdown(self.contents[title]) if isinstance(content, dict): self.display_structure(content, level + 1) def display_toc(self, structure, columns, level=1, col_index=0): for title, content in structure.items(): with columns[col_index % len(columns)]: st.markdown(f"{' ' * (level-1) * 2}- {title}") col_index += 1 if isinstance(content, dict): col_index = self.display_toc(content, columns, level + 1, col_index) return col_index def get_markdown_content(self, structure=None, level=1): """ Returns the markdown styled pure string with the contents. """ if structure is None: structure = self.structure if level==1: markdown_content = f"# {self.paper_title}\n\n" else: markdown_content = "" for title, content in structure.items(): if self.contents[title].strip(): # Only include title if there is content markdown_content += f"{'#' * level} {title}\n{self.contents[title]}\n\n" if isinstance(content, dict): markdown_content += self.get_markdown_content(content, level + 1) return markdown_content def create_markdown_file(content: str) -> BytesIO: """ Create a Markdown file from the provided content. """ markdown_file = BytesIO() markdown_file.write(content.encode("utf-8")) markdown_file.seek(0) return markdown_file def create_pdf_file(content: str) -> str: """ Create a PDF file from the provided Markdown content. Converts Markdown to styled HTML, then HTML to PDF. """ try: html_content = markdown(content, extensions=["extra", "codehilite"]) styled_html = f""" {html_content} """ with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file: HTML(string=styled_html).write_pdf(tmp_file.name) tmp_file.seek(0) return tmp_file.name except Exception as e: st.error(f"Error creating PDF file: {e}") return None def generate_paper_title(prompt: str, language: str): """ Generate a research paper title using AI. """ if language == "Arabic": prompt_language = "Arabic" else: prompt_language = "English" completion = st.session_state.groq.chat.completions.create( model="mixtral-8x7b-32768", messages=[ { "role": "system", "content": f"Generate suitable research paper titles for the provided topics in {prompt_language}. There is only one generated paper title! Don't give any explanation or add any symbols, just write the title of the paper. The requirement for this title is that it must be between 7 and 25 words long, and it must be attractive enough!" }, { "role": "user", "content": f"Generate a research paper title for the following topic in {prompt_language}. There is only one generated paper title! Don't give any explanation or add any symbols, just write the title of the paper. The requirement for this title is that it must be at least 7 words and 25 words long, and it must be attractive enough:\n\n{prompt}" } ], temperature=0.7, max_tokens=100, top_p=1, stream=False, stop=None, ) return completion.choices[0].message.content.strip() def generate_paper_structure(prompt: str, language: str): """ Returns research paper structure content as well as total tokens and total time for generation. """ if language == "Arabic": prompt_language = "Arabic" else: prompt_language = "English" completion = st.session_state.groq.chat.completions.create( model="mixtral-8x7b-32768", messages=[ { "role": "system", "content": f'Write in JSON format in {prompt_language}:\n\n{{"Title of section goes here":"Description of section goes here",\n"Title of section goes here":{{"Title of section goes here":"Description of section goes here","Title of section goes here":"Description of section goes here","Title of section goes here":"Description of section goes here"}}}}', }, { "role": "user", "content": f"Compose a detailed and comprehensive structure for an extensive research paper exceeding 300 pages. The structure should exclude sections such as the introduction and conclusion (including the foreword, author's note, and summary). The structure should be developed in {prompt_language} and should adhere closely to the following subject and additional instructions. Ensure the structure is well-organized and covers all necessary aspects of the topic comprehensively.\n\nالموضوع: {prompt}\n\nتعليمات إضافية: {additional_instructions}", } ], temperature=0.3, max_tokens=8000, top_p=1, stream=False, response_format={"type": "json_object"}, stop=None, ) usage = completion.usage statistics_to_return = GenerationStatistics( input_time=usage.prompt_time, output_time=usage.completion_time, input_tokens=usage.prompt_tokens, output_tokens=usage.completion_tokens, total_time=usage.total_time, model_name="mixtral-8x7b-32768", ) return statistics_to_return, completion.choices[0].message.content def generate_section(prompt: str, additional_instructions: str, language: str): if language == "Arabic": prompt_language = "Arabic" else: prompt_language = "English" stream = st.session_state.groq.chat.completions.create( model="mixtral-8x7b-32768", messages=[ { "role": "system", "content": f"You are an expert writer. Generate a long, comprehensive, structured chapter for the section provided in {prompt_language}. If additional instructions are provided, consider them very important. Only output the content.", }, { "role": "user", "content": f""" Generate a long, comprehensive, and well-structured chapter in {prompt_language}. Please adhere to the following guidelines: 1. **Section Title**: {prompt} 2. **Additional Instructions**: {additional_instructions} ### Guidelines: - Ensure the content is detailed and informative. - Maintain a logical flow and clear structure throughout the chapter. - Use appropriate headings and subheadings to organize the content. - Provide examples, case studies, or real-life applications where relevant. - Cite any external sources or references properly. - Keep the language formal and academic. ### Structure: - **Introduction**: Briefly introduce the topic and its importance. - **Main Body**: - Present key concepts and ideas. - Discuss various perspectives and arguments. - Include relevant data, statistics, and evidence. - **Conclusion**: Summarize the key points and provide any final insights or recommendations. Please make sure the content is engaging and free of grammatical errors. """ }, ], temperature=0.3, max_tokens=8000, top_p=1, stream=True, stop=None, ) for chunk in stream: tokens = chunk.choices[0].delta.content if tokens: yield tokens if x_groq := chunk.x_groq: if not x_groq.usage: continue usage = x_groq.usage statistics_to_return = GenerationStatistics( input_time=usage.prompt_time, output_time=usage.completion_time, input_tokens=usage.prompt_tokens, output_tokens=usage.completion_tokens, total_time=usage.total_time, model_name="mixtral-8x7b-32768", ) yield statistics_to_return def extract_text_from_pdf(file): """ Extracts text from the provided PDF file. """ document = fitz.open(stream=file.read(), filetype="pdf") text = "" for page_num in range(len(document)): page = document.load_page(page_num) text += page.get_text("text") return text def generate_research_citations(extracted_texts, language: str): """ Generate proper citations for the extracted texts. """ citations = [] for text in extracted_texts: if language == "Arabic": prompt_language = "Arabic" else: prompt_language = "English" citation = st.session_state.groq.chat.completions.create( model="mixtral-8x7b-32768", messages=[ { "role": "system", "content": f"Generate proper citations in APA format for the given extracted text from a research paper in {prompt_language}." }, { "role": "user", "content": f"Generate a citation for the following text in {prompt_language}:\n\n{text}" } ], temperature=0.7, max_tokens=100, top_p=1, stream=False, stop=None, ) citations.append(citation.choices[0].message.content.strip()) return citations def preprocess_texts(texts): """ Preprocess texts for indexing. """ preprocessed_texts = [] for text in texts: # Split text into smaller chunks for indexing chunks = text.split("\n\n") preprocessed_texts.extend(chunks) return preprocessed_texts def index_texts(texts): """ Index texts using FAISS. """ vectorizer = TfidfVectorizer() vectors = vectorizer.fit_transform(texts).toarray() index = faiss.IndexFlatL2(vectors.shape[1]) index.add(vectors) return index, vectorizer def retrieve_passages(query, index, vectorizer, texts, top_k=5): """ Retrieve top-k passages relevant to the query. """ query_vector = vectorizer.transform([query]).toarray() distances, indices = index.search(query_vector, top_k) retrieved_passages = [texts[i] for i in indices[0]] return retrieved_passages def split_text(text, max_tokens=1500): """ Splits text into smaller chunks to avoid exceeding the API's size limit. """ words = text.split() chunks = [] current_chunk = [] for word in words: current_chunk.append(word) if len(current_chunk) >= max_tokens: chunks.append(' '.join(current_chunk)) current_chunk = [] if current_chunk: chunks.append(' '.join(current_chunk)) return chunks # Initialize if "button_disabled" not in st.session_state: st.session_state.button_disabled = False if "button_text" not in st.session_state: st.session_state.button_text = "Generate" if "statistics_text" not in st.session_state: st.session_state.statistics_text = "" if 'paper_title' not in st.session_state: st.session_state.paper_title = "" if 'uploaded_pdfs' not in st.session_state: st.session_state.uploaded_pdfs = [] if 'extracted_texts' not in st.session_state: st.session_state.extracted_texts = [] if 'citations' not in st.session_state: st.session_state.citations = [] if 'index' not in st.session_state: st.session_state.index = None if 'vectorizer' not in st.session_state: st.session_state.vectorizer = None if 'preprocessed_texts' not in st.session_state: st.session_state.preprocessed_texts = [] st.write( """ # ResearchPaper: Write full research papers using AI """ ) def disable(): st.session_state.button_disabled = True def enable(): st.session_state.button_disabled = False def empty_st(): st.empty() try: if st.button("End Generation and Download Paper"): if "paper" in st.session_state: # Create markdown file markdown_file = create_markdown_file( st.session_state.paper.get_markdown_content() ) st.download_button( label="Download Text", data=markdown_file, file_name=f'{st.session_state.paper_title}.txt', mime='text/plain' ) # Create pdf file (styled) pdf_file_path = create_pdf_file(st.session_state.paper.get_markdown_content()) if pdf_file_path: with open(pdf_file_path, "rb") as pdf_file: pdf_data = pdf_file.read() st.download_button( label="Download PDF", data=pdf_data, file_name=f'{st.session_state.paper_title}.pdf', mime='application/pdf' ) else: st.error("Failed to generate the PDF file.") else: raise ValueError("Please generate content first before downloading the paper.") with st.form("groqform"): if not GROQ_API_KEY: groq_input_key = st.text_input( "Enter your Groq API Key (gsk_yA...):", "", type="password" ) topic_text = st.text_input( "What do you want the research paper to be about?", value="", help="Enter the main topic or title of your research paper", ) additional_instructions = st.text_area( "Additional Instructions (optional)", help="Provide any specific guidelines or preferences for the research paper's content", placeholder="E.g., 'Focus on beginner-friendly content', 'Include case studies', etc.", value="", ) # Language selection language = st.selectbox( "Choose the language for the research paper", options=["English", "Arabic"], index=0 ) # Upload PDFs uploaded_pdfs = st.file_uploader( "Upload related research PDFs", type=["pdf"], accept_multiple_files=True, help="Upload multiple PDFs that you want to extract information from and cite in the research paper", ) if uploaded_pdfs: for pdf in uploaded_pdfs: if pdf not in st.session_state.uploaded_pdfs: st.session_state.uploaded_pdfs.append(pdf) extracted_text = extract_text_from_pdf(pdf) st.session_state.extracted_texts.append(extracted_text) st.success(f"Extracted text from {pdf.name}") st.session_state.preprocessed_texts = preprocess_texts(st.session_state.extracted_texts) # Ensure there are non-stop words in the texts if len(st.session_state.preprocessed_texts) > 0: st.session_state.index, st.session_state.vectorizer = index_texts(st.session_state.preprocessed_texts) else: st.error("Uploaded documents do not contain enough content for processing. Please upload different documents.") # Generate button submitted = st.form_submit_button( st.session_state.button_text, on_click=disable, disabled=st.session_state.button_disabled, ) # Statistics placeholder = st.empty() def display_statistics(): with placeholder.container(): if st.session_state.statistics_text: if ( "Generating structure in background" not in st.session_state.statistics_text ): st.markdown( st.session_state.statistics_text + "\n\n---\n" ) # Format with line if showing statistics else: st.markdown(st.session_state.statistics_text) else: placeholder.empty() if submitted: if len(topic_text) < 10: raise ValueError("Research paper topic must be at least 10 characters long") st.session_state.button_disabled = True st.session_state.statistics_text = "Generating research paper title and structure in background...." display_statistics() if not GROQ_API_KEY: st.session_state.groq = Groq(api_key=groq_input_key) large_model_generation_statistics, paper_structure = generate_paper_structure( topic_text, language ) # Generate AI research paper title st.session_state.paper_title = generate_paper_title(topic_text, language) st.write(f"## {st.session_state.paper_title}") large_model_generation_statistics, paper_structure = generate_paper_structure(topic_text, language) total_generation_statistics = GenerationStatistics( model_name="mixtral-8x7b-32768" ) try: paper_structure_json = json.loads(paper_structure) paper = ResearchPaper(st.session_state.paper_title, paper_structure_json) if 'paper' not in st.session_state: st.session_state.paper = paper # Print the paper structure to the terminal to show structure print(json.dumps(paper_structure_json, indent=2)) st.session_state.paper.display_structure() def stream_section_content(sections): for title, content in sections.items(): if isinstance(content, str): if st.session_state.index and st.session_state.vectorizer: # Retrieve passages related to the section retrieved_passages = retrieve_passages( title + ": " + content, st.session_state.index, st.session_state.vectorizer, st.session_state.preprocessed_texts ) context = "\n\n".join(retrieved_passages) prompt_with_context = title + ": " + content + "\n\n" + context else: prompt_with_context = title + ": " + content # Split the prompt into smaller chunks chunks = split_text(prompt_with_context) for chunk in chunks: content_stream = generate_section( chunk, additional_instructions, language ) for content_chunk in content_stream: # Check if GenerationStatistics data is returned instead of str tokens if isinstance(content_chunk, GenerationStatistics): total_generation_statistics.add(content_chunk) st.session_state.statistics_text = str( total_generation_statistics ) display_statistics() elif content_chunk: st.session_state.paper.update_content(title, content_chunk) elif isinstance(content, dict): stream_section_content(content) stream_section_content(paper_structure_json) # Append extracted texts with citations if st.session_state.extracted_texts: citations = generate_research_citations(st.session_state.extracted_texts, language) st.session_state.citations = citations st.session_state.paper.update_content( "References", "\n\n".join(st.session_state.citations) ) except json.JSONDecodeError: st.error("Failed to decode the research paper structure. Please try again.") enable() except Exception as e: st.session_state.button_disabled = False st.error(e) if st.button("Clear"): st.rerun()