import gradio as gr import pandas as pd import numpy as np from langchain_openai import OpenAIEmbeddings, ChatOpenAI from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.vectorstores import FAISS from langchain.prompts import ChatPromptTemplate from langchain.schema.runnable import RunnablePassthrough from langchain.schema.output_parser import StrOutputParser import os from dotenv import load_dotenv # Load environment variables load_dotenv() OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") # Load CSV data df = pd.read_csv("K-Funding250331.csv") # Clean and prepare data def clean_and_prepare_data(df): # Fill NaN values with empty strings df_clean = df.fillna("") # Create document-like entries for each row documents = [] for idx, row in df_clean.iterrows(): # Convert the row to a string representation doc_content = f"Startup: {row['Startup Name']}\n" for col in df_clean.columns: if col != 'Startup Name' and row[col]: # Add only non-empty fields doc_content += f"{col}: {row[col]}\n" documents.append({"content": doc_content, "source": f"row_{idx}"}) return documents # Convert DataFrame to documents documents = clean_and_prepare_data(df) # Split documents into chunks text_splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=100 ) splits = text_splitter.create_documents([doc["content"] for doc in documents]) # Add metadata back to splits for i, split in enumerate(splits): split.metadata = {"source": f"chunk_{i}"} # Create vector store embeddings = OpenAIEmbeddings() vectorstore = FAISS.from_documents(splits, embeddings) retriever = vectorstore.as_retriever(search_kwargs={"k": 5}) # Function to translate both company and investor names from English to Korean def translate_entity_names(query): # Create a translation model translation_model = ChatOpenAI(model="gpt-4o-mini", temperature=0) # Updated prompt to handle both company and investor names translation_prompt = ChatPromptTemplate.from_messages([ ("system", """You are a Korean-English translator specializing in Korean startup and investor names. Extract any potential company or investor names from the user's query and provide their Korean translations. If you detect company or investor names, return them in this format: "TRANSLATION_RESULT: {{English Name}} -> {{Korean Name}} (Type: company/investor)" You can return multiple translations if multiple entities are detected: "TRANSLATION_RESULT: {{English Name 1}} -> {{Korean Name 1}} (Type: company) {{English Name 2}} -> {{Korean Name 2}} (Type: investor)" If no relevant names are detected, return: "TRANSLATION_RESULT: NO_ENTITIES_DETECTED" Do not provide any explanations or additional text, just the translation result."""), ("human", "{query}") ]) translation_chain = translation_prompt | translation_model | StrOutputParser() result = translation_chain.invoke({"query": query}) # Process the translation result translations = [] modified_query = query if "TRANSLATION_RESULT:" in result: result_parts = result.strip().split("TRANSLATION_RESULT:")[1].strip() if result_parts != "NO_ENTITIES_DETECTED": # Handle multiple translations if present translation_lines = [line.strip() for line in result_parts.split('\n') if line.strip()] for line in translation_lines: if "->" in line: # Extract parts name_part, type_part = line.rsplit("(Type:", 1) if "(Type:" in line else (line, "(Type: unknown)") english_name, korean_name = name_part.split("->") entity_type = type_part.replace(")", "").strip() translation = { "english_name": english_name.strip(), "korean_name": korean_name.strip(), "entity_type": entity_type.strip() } translations.append(translation) # Replace the English name with Korean name in the query modified_query = modified_query.replace(english_name.strip(), korean_name.strip()) if translations: return { "translated": True, "translations": translations, "modified_query": modified_query } else: return { "translated": False, "original_query": query } # Create a system template that instructs the model how to handle the data system_template = """You are an expert financial analyst specializing in Korean startup investments. Answer the user's question based on the following retrieved information. If you don't know the answer, say that you don't know. Use the retrieved information to provide insights, summaries, and analysis. Retrieved information: {context} Additional instructions: - Translate any Korean text to English in your response - If comparing multiple startups or sectors, create a structured comparison - If analyzing trends, provide clear patterns and insights - Be precise with numbers and statistics - If information is missing, acknowledge it rather than making assumptions - When referencing any company name, use both the English translated name and also the original Korean name in bracket - Do the same for investor names - provide both English and Korean names where possible """ # Create a human template for the user's questions human_template = "{question}" # Create a chat prompt template chat_prompt = ChatPromptTemplate.from_messages([ ("system", system_template), ("human", human_template), ]) # Create a chat model model = ChatOpenAI(model="gpt-4o-mini", temperature=0) # Define the RAG chain def retrieve_and_generate(query_info): # Extract the query and context from the input query = query_info["question"] context = query_info.get("context", "") # Run the chat model result = chat_prompt.invoke({ "context": context, "question": query }) return model.invoke(result) # Function to handle specialized queries def handle_query(query): # First, attempt to translate any company or investor names in the query translation_result = translate_entity_names(query) # If translation was successful, use the modified query for retrieval if translation_result.get("translated", False): original_query = query query_for_retrieval = translation_result["modified_query"] # Add translation info to the query for the model translation_notes = [] for t in translation_result["translations"]: translation_notes.append(f"'{t['english_name']}' to '{t['korean_name']}' (Type: {t['entity_type']})") translation_info = ", ".join(translation_notes) enhanced_query = f"{original_query}\n\nNote: I've translated: {translation_info} for better retrieval." else: query_for_retrieval = query enhanced_query = query # Check if query is about overall statistics if any(word in query.lower() for word in ["all", "overall", "summary", "statistics", "stats"]): # Generate statistical summary summary_stats = generate_statistics(df) # Add the statistics to the context context = f"Overall Dataset Statistics:\n{summary_stats}\n\n" # Use the retriever with the translated query retrieved_docs = retriever.invoke(query_for_retrieval) retrieved_context = "\n\n".join([doc.page_content for doc in retrieved_docs]) combined_context = context + retrieved_context # Run the query with the statistics and retrieved context response = retrieve_and_generate({ "context": combined_context, "question": enhanced_query }) return StrOutputParser().invoke(response) # For all other queries, use the retriever with the translated query retrieved_docs = retriever.invoke(query_for_retrieval) retrieved_context = "\n\n".join([doc.page_content for doc in retrieved_docs]) # Run the query with the retrieved context response = retrieve_and_generate({ "context": retrieved_context, "question": enhanced_query }) return StrOutputParser().invoke(response) # Function to generate statistics from the DataFrame def generate_statistics(df): stats = [] # Count total number of startups stats.append(f"Total startups: {len(df)}") # Count unique sectors sectors = df["Sector"].dropna().unique() stats.append(f"Unique sectors: {len(sectors)}") # Count by funding rounds if "Round (KR)" in df.columns: round_counts = df["Round (KR)"].value_counts().to_dict() stats.append("Funding round distribution:") for round_type, count in round_counts.items(): stats.append(f" - {round_type}: {count}") # Calculate average investment if available if "Investment Amount" in df.columns: # This is simplified, need to properly parse the investment amounts stats.append("Investment statistics available in the dataset") return "\n".join(stats) # Define the Gradio interface demo = gr.Interface( fn=handle_query, inputs=gr.Textbox(label="Ask a question about Korean startup investments"), outputs=gr.Textbox(label="AI Answer"), title="Korean Startup Knowledge Base", description="Ask questions about Korean startups, funding trends, sectors, or investment rounds.", examples=[ ["Which investors are most active in the healthcare sector?"], ["What was the funding trend for AI startups in 2025 so far?"], ["What do you know about Spoon Labs?"], ["Tell me about investments made by Kakao Investment."] ] ) # Launch the Gradio interface if __name__ == "__main__": demo.launch()