from saged import Pipeline from tqdm import tqdm from pathlib import Path from saged import SAGEDData as dt import streamlit as st import json import http.client from openai import AzureOpenAI import ollama import time # Use time.sleep to simulate processing steps import logging from io import StringIO import sys # Create a custom logging handler to capture log messages class StreamlitLogHandler(logging.Handler): def __init__(self): super().__init__() self.log_capture_string = StringIO() def emit(self, record): # Write each log message to the StringIO buffer message = self.format(record) self.log_capture_string.write(message + "\n") def get_logs(self): # Return the log contents return self.log_capture_string.getvalue() def clear_logs(self): # Clear the log buffer self.log_capture_string.truncate(0) self.log_capture_string.seek(0) # Define ContentFormatter class class ContentFormatter: @staticmethod def chat_completions(text, settings_params): message = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": text} ] data = {"messages": message, **settings_params} return json.dumps(data) # Define OllamaModel (For local Ollama interaction) class OllamaModel: def __init__(self, base_model='llama3', system_prompt='You are a helpful assistant', model_name='llama3o', **kwargs): self.base_model = base_model self.model_name = model_name self.model_create(model_name, system_prompt, base_model, **kwargs) def model_create(self, model_name, system_prompt, base_model, **kwargs): modelfile = f'FROM {base_model}\nSYSTEM {system_prompt}\n' if kwargs: for key, value in kwargs.items(): modelfile += f'PARAMETER {key.lower()} {value}\n' ollama.create(model=model_name, modelfile=modelfile) def invoke(self, prompt): answer = ollama.generate(model=self.model_name, prompt=prompt) return answer['response'] # Define GPTAgent (For OpenAI GPT models) class GPTAgent: def __init__(self, model_name, azure_key, azure_version, azure_endpoint, deployment_name): self.client = AzureOpenAI( api_key=azure_key, api_version=azure_version, azure_endpoint=azure_endpoint ) self.deployment_name = deployment_name def invoke(self, prompt, settings_params=None): if not settings_params: settings_params = {} formatted_input = ContentFormatter.chat_completions(prompt, settings_params) response = self.client.chat.completions.create( model=self.deployment_name, messages=json.loads(formatted_input)['messages'], **settings_params ) return response.choices[0].message.content # Define AzureAgent (For Azure OpenAI models) class AzureAgent: def __init__(self, model_name, azure_uri, azure_api_key): self.azure_uri = azure_uri self.headers = { 'Authorization': f"Bearer {azure_api_key}", 'Content-Type': 'application/json' } self.chat_formatter = ContentFormatter def invoke(self, prompt, settings_params=None): if not settings_params: settings_params = {} body = self.chat_formatter.chat_completions(prompt, {**settings_params}) conn = http.client.HTTPSConnection(self.azure_uri) conn.request("POST", '/v1/chat/completions', body=body, headers=self.headers) response = conn.getresponse() data = response.read() conn.close() decoded_data = data.decode("utf-8") parsed_data = json.loads(decoded_data) content = parsed_data["choices"][0]["message"]["content"] return content # Renew Source Finder Button def renew_source_finder(domain, concept_list): if 'generated_synthetic_files' in st.session_state: del st.session_state['generated_synthetic_files'] if not domain or not concept_list: st.error("Please fill in all the required fields before proceeding.") else: with st.spinner("Renewing source info files..."): base_path = Path('data/customized/source_finder/') for concept in concept_list: file_path = base_path / f'{domain}_{concept}_source_finder.json' if file_path.exists(): try: file_path.unlink() # Delete the file st.info(f"Deleted source info file: {file_path}") except Exception as e: st.error(f"An error occurred while deleting the file {file_path}: {e}") st.success("Source info files renewal completed!") def create_source_finder(domain, concept): source_specification_item = f"data/customized/local_files/{domain}/{concept}.txt" if not Path(source_specification_item).exists(): st.warning(f"Local file does not exist: {source_specification_item}") instance = dt.create_data(domain, concept, 'source_finder') instance.data[0]['keywords'] = {concept: dt.default_keyword_metadata.copy()} category_shared_source_item = dt.default_source_item.copy() category_shared_source_item['source_type'] = "local_paths" category_shared_source_item['source_specification'] = [source_specification_item] instance.data[0]['category_shared_source'] = [category_shared_source_item] return instance.data.copy() def check_and_create_source_files(domain, concept_list): """ Checks if the required source finder files exist for each concept in the domain. If a file does not exist or is invalid, it creates an empty JSON file for that concept. """ base_path = Path('data/customized/source_finder/') base_path.mkdir(parents=True, exist_ok=True) for concept in concept_list: file_path = base_path / f'{domain}_{concept}_source_finder.json' if not file_path.exists(): # Create a new source finder file using create_source_finder data = create_source_finder(domain, concept) with open(file_path, 'w', encoding='utf-8') as f: json.dump(data, f, indent=4) st.info(f"Created missing source finder file: {file_path}") else: # Attempt to load the file to verify its validity instance = dt.load_file(domain, concept, 'source_finder', file_path) if instance is None: # If loading fails, create a new valid file data = create_source_finder(domain, concept) with open(file_path, 'w', encoding='utf-8') as f: json.dump(data, f, indent=4) st.info(f"Recreated invalid source finder file: {file_path}") def clean_spaces(data): """ Removes trailing or leading spaces from a string or from each element in a list. """ if isinstance(data, str): return data.strip() elif isinstance(data, list): return [item.strip() if isinstance(item, str) else item for item in data] else: raise TypeError("Input should be either a string or a list of strings") def create_replacement_dict(concept_list, replacer): replacement = {} for concept in concept_list: replacement[concept] = {} for company in replacer: replacement[concept][company] = {concept: company} return replacement # Title of the app st.title("SAGED-bias Benchmark-Building Demo") # Initialize session state variables if 'domain' not in st.session_state: st.session_state['domain'] = None if 'concept_list' not in st.session_state: st.session_state['concept_list'] = None if 'gpt_model' not in st.session_state: st.session_state['gpt_model'] = None if 'azure_model' not in st.session_state: st.session_state['azure_model'] = None if 'ollama_model' not in st.session_state: st.session_state['ollama_model'] = None # Sidebar: Model Selection with st.sidebar: st.header("Model Configuration") # Selection of which model to use model_selection = st.radio("Select Model Type", ['GPT-Azure', 'Azure', 'Ollama']) # Collapsible Additional Configuration Section with st.expander("Model Configuration"): if model_selection == 'Ollama': # Ollama Configuration ollama_deployment_name = st.text_input("Enter Ollama Model Deployment Name", placeholder="e.g., llama3") ollama_system_prompt = st.text_input("Enter System Prompt for Ollama", placeholder="e.g., You are a helpful assistant.") if ollama_deployment_name and ollama_system_prompt: confirm_ollama = st.button("Confirm Ollama Configuration") if confirm_ollama: st.session_state['ollama_model'] = OllamaModel( model_name=ollama_deployment_name, system_prompt=ollama_system_prompt ) st.success("Ollama model configured successfully.") else: st.warning("Please provide both Ollama deployment name and system prompt.") elif model_selection == 'GPT-Azure' or model_selection == 'Azure': # GPT / Azure Configuration gpt_azure_endpoint = st.text_input("Enter Azure Endpoint URL", placeholder="e.g., https://your-resource-name.openai.azure.com/") gpt_azure_api_key = st.text_input("Enter Azure API Key", type="password") gpt_azure_model_name = st.text_input("Enter Azure Model Name", placeholder="e.g., GPT-3.5-turbo") gpt_azure_deployment_name = st.text_input("Enter Azure Deployment Name", placeholder="e.g., gpt-3-5-deployment") if gpt_azure_endpoint and gpt_azure_api_key and gpt_azure_model_name and gpt_azure_deployment_name: confirm_gpt_azure = st.button("Confirm GPT/Azure Configuration") if confirm_gpt_azure: if model_selection == 'GPT-Azure': st.session_state['gpt_model'] = GPTAgent( model_name=gpt_azure_model_name, azure_key=gpt_azure_api_key, azure_version='2023-05-15', # Update if necessary azure_endpoint=gpt_azure_endpoint, deployment_name=gpt_azure_deployment_name ) st.success("GPT model configured successfully.") elif model_selection == 'Azure': st.session_state['azure_model'] = AzureAgent( model_name=gpt_azure_model_name, azure_uri=gpt_azure_endpoint, azure_api_key=gpt_azure_api_key ) st.success("Azure model configured successfully.") else: st.warning("Please provide all fields for GPT/Azure configuration.") # Main interaction based on configured model if st.session_state.get('ollama_model'): model = st.session_state['ollama_model'] elif st.session_state.get('gpt_model'): model = st.session_state['gpt_model'] elif st.session_state.get('azure_model'): model = st.session_state['azure_model'] else: model = None # User input: Domain and Concepts with st.form(key='domain_concept_form'): domain = clean_spaces( st.text_input("Enter the domain: (e.g., Stocks, Education)", placeholder="Enter domain here...")) # User input: Concepts concept_text = st.text_area("Enter the concepts (separated by commas):", placeholder="e.g., excel-stock, ok-stock, bad-stock") concept_list = clean_spaces(concept_text.split(',')) submit_button = st.form_submit_button(label='Confirm Domain and Concepts') if submit_button: if not domain: st.warning("Please enter a domain.") elif not concept_list or concept_text.strip() == "": st.warning("Please enter at least one concept.") else: st.session_state['domain'] = domain st.session_state['concept_list'] = concept_list st.success("Domain and concepts confirmed.") # Display further options only after domain and concepts are confirmed if st.session_state['domain'] and st.session_state['concept_list']: with st.expander("Additional Options"): # User input: Method scraper_method = st.radio("Select the scraper method:", (('wiki', 'local_files', 'synthetic_files'))) # Initiate the source_finder_requirement and keyword_finder_requirement if 'wiki' is selected if scraper_method == 'wiki': st.session_state['keyword_finder_requirement'] = True st.session_state['source_finder_requirement'] = True st.session_state['check_source_finder'] = False # File upload for each concept if 'local_files' is selected if scraper_method == 'local_files': uploaded_files = {} st.session_state['keyword_finder_requirement'] = False st.session_state['source_finder_requirement'] = False st.session_state['check_source_finder'] = True for concept in st.session_state['concept_list']: uploaded_file = st.file_uploader(f"Upload file for concept '{concept}':", type=['txt'], key=f"file_{concept}") if uploaded_file: uploaded_files[concept] = uploaded_file # Save uploaded file save_path = Path(f"data/customized/local_files/{st.session_state['domain']}/{concept}.txt") save_path.parent.mkdir(parents=True, exist_ok=True) with open(save_path, 'wb') as f: f.write(uploaded_file.getbuffer()) st.success(f"File for concept '{concept}' saved successfully.") # Generate synthetic files if 'synthetic_files' is selected if scraper_method == 'synthetic_files': scraper_method = 'local_files' st.session_state['keyword_finder_requirement'] = False st.session_state['source_finder_requirement'] = False st.session_state['check_source_finder'] = True if 'generated_synthetic_files' not in st.session_state: st.session_state['generated_synthetic_files'] = set() prompt_inputs = {} for concept in st.session_state['concept_list']: if concept not in st.session_state['generated_synthetic_files']: prompt_inputs[concept] = st.text_input( f"Enter the prompt for concept '{concept}':", value=f"Write a long article introducing the {concept} in the {st.session_state['domain']}. Use the {concept} as much as possible.", key=f"prompt_{concept}" ) if st.button("Generate Synthetic Files for All Concepts"): if model: for concept, prompt in prompt_inputs.items(): if prompt: with st.spinner(f"Generating content for concept '{concept}'..."): synthetic_content = model.invoke(prompt) save_path = Path( f"data/customized/local_files/{st.session_state['domain']}/{concept}.txt") save_path.parent.mkdir(parents=True, exist_ok=True) with open(save_path, 'w', encoding='utf-8') as f: f.write(synthetic_content) st.session_state['generated_synthetic_files'].add(concept) st.success(f"Synthetic file for concept '{concept}' created successfully.") else: st.warning("Please configure a model to generate synthetic files.") # User input: Prompt Method prompt_method = st.radio("Select the prompt method:", ('split_sentences', 'questions'), index = 0) # User input: Max Benchmark Length max_benchmark_length = st.slider("Select the maximum prompts per concepts:", 1, 199, 10) # User input: Branching branching = st.radio("Enable branching:", ('Yes', 'No'), index=1) branching_enabled = True if branching == 'Yes' else False # User input: Replacer (only if branching is enabled) replacer = [] replacement = {} if branching_enabled: replacer_text = st.text_area("Enter the replacer list (list of strings, separated by commas):", placeholder="e.g., Company A, Company B") replacer = clean_spaces(replacer_text.split(',')) replacement = create_replacement_dict(st.session_state['concept_list'], replacer) # Configuration concept_specified_config = { x: {'keyword_finder': {'manual_keywords': [x]}} for x in st.session_state['concept_list'] } concept_configuration = { 'keyword_finder': { 'require': st.session_state['keyword_finder_requirement'], 'keyword_number': 1, }, 'source_finder': { 'require': st.session_state['source_finder_requirement'], 'scrap_number': 10, 'method': scraper_method, }, 'scraper': { 'require': True, 'method': scraper_method, }, 'prompt_maker': { 'method': prompt_method, 'generation_function': model.invoke if model else None, 'max_benchmark_length': max_benchmark_length, }, } domain_configuration = { 'categories': st.session_state['concept_list'], 'branching': branching_enabled, 'branching_config': { 'generation_function': model.invoke if model else None, 'keyword_reference': st.session_state['concept_list'], 'replacement_descriptor_require': False, 'replacement_description': replacement, 'branching_pairs': 'not all', 'direction': 'not both', }, 'shared_config': concept_configuration, 'category_specified_config': concept_specified_config } # Renew Source Finder Button if st.button('Renew Source info'): renew_source_finder(st.session_state['domain'], st.session_state['concept_list']) # Save the original stdout to print to the terminal if needed later original_stdout = sys.stdout # Define StreamToText to capture and display logs in real-time within Streamlit only class StreamToText: def __init__(self): self.output = StringIO() def write(self, message): if message.strip(): # Avoid adding empty messages # Only append to Streamlit display, not the terminal st.session_state.log_messages.append(message.strip()) log_placeholder.text("\n".join(st.session_state.log_messages)) # Flush updated logs def flush(self): pass # Required for compatibility with sys.stdout # Initialize session state for log messages if 'log_messages' not in st.session_state: st.session_state.log_messages = [] # Replace sys.stdout with our custom StreamToText instance stream_to_text = StreamToText() sys.stdout = stream_to_text # Placeholder for displaying logs within a collapsible expander with st.expander("Show Logs", expanded=False): log_placeholder = st.empty() # Placeholder for dynamic log display # Define the Create Benchmark button if st.button("Create a Benchmark"): st.session_state.log_messages = [] # Clear previous logs with st.spinner("Creating benchmark..."): if st.session_state['check_source_finder']: # Check for relevant materials check_and_create_source_files(st.session_state['domain'], st.session_state['concept_list']) try: # Display progress bar and log messages progress_bar = st.progress(0) for i in tqdm(range(1, 101)): progress_bar.progress(i) time.sleep(0.05) # Short delay to simulate processing time # Run the benchmark creation function benchmark = Pipeline.domain_benchmark_building(st.session_state['domain'], domain_configuration) st.success("Benchmark creation completed!") st.dataframe(benchmark.data) except Exception as e: st.error(f"An error occurred during benchmark creation: {e}")