# -*- coding: utf-8 -*-
import gradio as gr
import pandas as pd
import json
import random
from LLM import zero_shot
from prompt_generate import generate_prompt_with_examples as generate_prompt
from prompt_generate import generate_prompt_with_best_matches as generate_prompt_b

def get_model_options():
    """Get available model series options"""
    return ['gpt', 'llama', 'qwen', 'deepSeek', 'gemini', 'claude']

def get_common_model_names(model_series):
    """Return common model name options based on model series"""
    model_names = {
        'gpt': ['gpt-3.5-turbo', 'gpt-4o'],
        'llama': ['meta-llama/Meta-Llama-3.1-405B-Instruct'],
        'qwen': ['Qwen/Qwen2.5-72B-Instruct'],
        'deepSeek': ['deepseek-ai/DeepSeek-V3', 'deepseek-ai/DeepSeek-R1'],
        'gemini': ['gemini-1.5-pro-002'],
        'claude': ['claude-3-5-haiku-20241022']
    }
    return model_names.get(model_series, [])

def get_prompt_templates():
    """Get predefined prompt templates"""
    templates = {
        "Custom": "",
        "Zero-shot Basic Extraction": """You are a professional and experienced expert in engineering geology. Your task is to extract "entity-relation-entity" triples from the given input text. There are 24 types of relations: "outcrop at", "located in", "conformable contact", "unconformable contact", "paraconformable contact", "fault contact", "distribution pattern", "tectonic position", "stratigraphic division", "exposed strata", "lithology", "thickness", "area", "coordinates", "length", "contains", "age", "administrative division", "development", "paleontology", "elevation", "belongs to", "engulf", "intrude". Please follow these specifications for extraction:
1. Output format:
Strictly follow JSON array format, no additional text, each element contains:
[
    {
        "entity1": "Entity1",
        "relation": "Relation",
        "entity2": "Entity2"
    }
]
2. Complex relationship handling:
- If the same entity participates in multiple relationships, list different triples separately""",
        
        "knowledge-injected Enhanced Extraction": """You are a professional and experienced expert in engineering geology. Your task is to extract "entity-relation-entity" triples from the given input text. There are 24 types of relations: "outcrop at", "located in", "conformable contact", "unconformable contact", "paraconformable contact", "fault contact", "distribution pattern", "tectonic position", "stratigraphic division", "exposed strata", "lithology", "thickness", "area", "coordinates", "length", "contains", "age", "administrative division", "development", "paleontology", "elevation", "belongs to", "engulf", "intrude". Please follow these specifications for extraction:
1. Output format:
Strictly follow JSON array format, no additional text, each element contains:
[
    {
        "entity1": "Entity1",
        "relation": "Relation",
        "entity2": "Entity2"
    }
]
2. Complex relationship handling:
- If the same entity participates in multiple relationships, list different triples separately
3. Relationship explanations:
Outcrop at: Refers to rocks or strata exposed at the surface or near-surface, not covered or buried. Example: (Late Ordovician-Silurian intrusive rocks, outcrop at, southern investigation area).

Located in: Establishes the subordinate relationship of geological units within a larger spatial framework (administrative region/tectonic unit). Example: (Kumuqi Silurian basaltic basic rocks, located in, central-western investigation area)

Conformable contact: Indicates contact relationships formed by continuous deposition of upper and lower strata, reflecting gradational lithological characteristics without significant depositional hiatus. Example: (Solake Formation, conformable contact, Middle Ordovician Lin Formation).

Unconformable contact: Describes stratigraphic contact interfaces with depositional gaps, including contact features with angular differences or lithological abrupt changes. Example: (Tongziyan Formation, unconformable contact, Maokou Formation).

Paraconformable contact: Specifically refers to parallel unconformity types with consistent attitudes, emphasizing depositional sequence interruption but without structural deformation. Example: (Solake Formation, paraconformable contact, Middle Ordovician Lin Formation).

Fault contact: Two strata are separated by fault zones or fault planes, often accompanied by dynamic crushing and other structural phenomena. Example: (Solake Formation, fault contact, Upper Ordovician Lapai Spring Formation).

Distribution pattern: Depicts spatial distribution characteristics of geological units, including geometric morphology and extension direction combinations. Example: (Carboniferous, distribution pattern, banded).

Tectonic position: Locates geological units' attribution in plate tectonic framework, associated with orogenic belts or tectonic unit divisions. Example: (Carboniferous, tectonic position, northern margin of Gondwana tectonic belt).

Stratigraphic division: Characterizes hierarchical attribution and zoning attributes of stratigraphic units in regional stratigraphic division systems. Example: (Carboniferous, stratigraphic division, Gondwana).

Exposed strata: Specifically refers to actually exposed stratigraphic entities in a region, emphasizing observable surface geological units. Example: (Hongliugou gold-copper mining area, exposed strata, Nanhua-Lower Ordovician Hongliugou Group).

Lithology: Defines material composition and structural characteristics of rocks, including hierarchical descriptive elements of composite lithology. Example: (Late Ordovician-Silurian syenite, lithology, altered syenite).

Thickness: Quantifies vertical dimensions of strata/rock bodies, including dimensional expressions with absolute values and relative descriptions. Example: (syenite, thickness, 35.60 m).

Area: Characterizes horizontal distribution range of geological units, presented in standardized form combining numerical values and units. Example: (intrusive rocks, outcrop area, 54 m2)

Coordinates: Specifically refers to geographical spatial positioning data recording geological feature points. Example: (Solake copper-gold mine site, coordinates, 90°11′47″E).

Length: Describes spatial extension dimensions of linear geological bodies. Example: Triple (Shibien fault zone, length, 20m) can be extracted.

Contains: Indicates compositional inclusion relationships of main materials, specifically referring to mineral composition or fossil occurrence states, different from everyday meaning. Example: (medium gray-black massive chert, contains, chert bands).

Age: Establishes correspondence between geological units and standard geological chronological systems. Example: (Hongliugou gold-copper mining area, age, Early-Middle Permian).

Administrative division: Defines subordinate hierarchy and territorial attribution of geological entities in administrative management systems. Example: (investigation area, administrative division, Chayang County).

Development: Describes manifestation degree and formation state intensity of geological structures or depositional features. Example: (Lanhuaweng Formation, development, horizontal bedding).

Paleontology: Records fossil biological information occurring in strata, requiring complete Latin scientific names and classification features. Example: (strata, paleontology, Lumu et al).

Elevation: Quantifies elevation data of geological feature points relative to sea level, retaining measurement reference identification. Example: (Solake copper-gold mine site, elevation, 2800m).

Belongs to: Establishes type attribution of geological units in classification systems. Example: (mining area, belongs to, polymetallic mineralization subarea).

Engulf: Characterizes spatial replacement processes of intrusive bodies on country rocks, reflecting transformation effects of magmatic activities. Example: (Nintendo Rock Formation, engulf, Jurassic granite).

Intrude: Describes geological processes of magmatic rock bodies penetrating country rocks, including accompanying phenomena such as contact metamorphism. Example: (Gaozhou Shell Stone Formation, intrude, gneissic granite).

4. Other key points:
All triple relationships must be one of the above 24 types
Relationship entities cannot be verbs, prepositions, or other meaningless words. Descriptions of rocks, strata, and other entities should be as complete as possible according to the original text""",
    }
    return templates

def get_qa_prompt_templates():
    """Get QA module prompt templates"""
    templates = {
        "Custom": "",
        "Zero-shot True/False": "Please judge true or false based on the given text.",
        "Zero-shot Q&A": "Please answer the question based on the given text.",
        "COT True/False": "Please first judge true or false, and provide your reasoning basis.",
        "COT Q&A": "Please first answer the question, and provide your reasoning basis.",
    }
    return templates

# Global variables to store training data
_train_data = None
_text_series = None
_label_series = None

def load_train_data():
    """Load training data"""
    global _train_data, _text_series, _label_series
    if _train_data is None:
        try:
            _train_data = pd.read_json('./data/train_triples.json')
            _text_series = _train_data['text']
            _label_series = _train_data['triple_list']
        except Exception as e:
            print(f"Failed to load training data: {e}")
            return False
    return True

def generate_random_context_prompt(user_text, num_examples):
    """Generate random context prompts"""
    if not load_train_data():
        return "Unable to load training data"
    
    try:
        random_prompt = generate_prompt(_text_series, _label_series, num_examples)
        return f"Here are geological description text and triple extraction examples:\n\n{random_prompt}\nPlease extract triples based on the examples:\n{user_text}"
    except Exception as e:
        return f"Failed to generate random context prompt: {e}"

def generate_best_match_context_prompt(user_text, num_examples):
    """Generate best match context prompts based on similarity"""
    if not load_train_data():
        return "Unable to load training data"
    
    try:
        best_match_prompt = generate_prompt_b(_text_series, _label_series, user_text, num_examples)
        if best_match_prompt.strip():
            return f"Here are geological description text and triple extraction examples:\n\n{best_match_prompt}\n\nPlease extract triples based on the examples:\n{user_text}"
        else:
            return f"No matching examples found, performing zero-shot extraction:\n{user_text}"
    except Exception as e:
        return f"Failed to generate best match context prompt: {e}"

def update_model_names(model_series):
    """Update model name dropdown list when model series changes"""
    names = get_common_model_names(model_series)
    return gr.Dropdown(choices=names, value=names[0] if names else "", label="Model Name", allow_custom_value=True)

def update_prompt_content(template_name):
    """Update content when prompt template changes"""
    templates = get_prompt_templates()
    content = templates.get(template_name, "")
    return gr.Textbox(value=content, label="Prompt Content", lines=15, max_lines=25)

def update_qa_prompt_content(template_name):
    """Update content when QA prompt template changes"""
    templates = get_qa_prompt_templates()
    content = templates.get(template_name, "")
    return gr.Textbox(value=content, label="QA Prompt Content", lines=3, max_lines=10)

def call_llm_model(model_series, model_name, prompt_content, user_content, context_type, num_examples):
    """LLM model wrapper function (triple extraction)"""
    try:
        if not model_series or not model_name:
            return "Please select model series and model name"
        
        if not user_content:
            return "Please input text content to process"
        
        # Combine complete input content based on context type
        if context_type == "No Context":
            if prompt_content.strip():
                full_content = prompt_content.strip() + "\n\n" + user_content
            else:
                full_content = user_content
        elif context_type == "Random Context":
            context_prompt = generate_random_context_prompt(user_content, num_examples)
            if prompt_content.strip():
                full_content = prompt_content.strip() + "\n\n" + context_prompt
            else:
                full_content = context_prompt
        elif context_type == "Best Match Context":
            context_prompt = generate_best_match_context_prompt(user_content, num_examples)
            if prompt_content.strip():
                full_content = prompt_content.strip() + "\n\n" + context_prompt
            else:
                full_content = context_prompt
        else:
            if prompt_content.strip():
                full_content = prompt_content.strip() + "\n\n" + user_content
            else:
                full_content = user_content
        
        response = zero_shot(model_series, model_name, full_content)
        
        # Handle different types of return values
        if hasattr(response, 'content'):
            return response.content
        elif isinstance(response, dict) and 'content' in response:
            return response['content']
        elif isinstance(response, str):
            return response
        else:
            return str(response)
            
    except Exception as e:
        return f"Error calling model: {str(e)}"

def call_qa_model(model_series, model_name, qa_prompt_content, geological_text, question_or_statement, qa_type):
    """LLM model wrapper function (QA module)"""
    try:
        if not model_series or not model_name:
            return "Please select model series and model name"
        
        if not geological_text:
            return "Please input geological text"
        
        if not question_or_statement:
            if qa_type == "True/False":
                return "Please input factual statement to judge"
            else:
                return "Please input question to answer"
        
        # Combine complete input content
        if qa_type == "True/False":
            if qa_prompt_content.strip():
                full_content = f"{qa_prompt_content.strip()}\n\nGeological text:\n{geological_text}\n\nStatement to judge:\n{question_or_statement}"
            else:
                full_content = f"Geological text:\n{geological_text}\n\nStatement to judge:\n{question_or_statement}"
        else:  # Q&A
            if qa_prompt_content.strip():
                full_content = f"{qa_prompt_content.strip()}\n\nGeological text:\n{geological_text}\n\nQuestion:\n{question_or_statement}"
            else:
                full_content = f"Geological text:\n{geological_text}\n\nQuestion:\n{question_or_statement}"
        
        response = zero_shot(model_series, model_name, full_content)
        
        # Handle different types of return values
        if hasattr(response, 'content'):
            return response.content
        elif isinstance(response, dict) and 'content' in response:
            return response['content']
        elif isinstance(response, str):
            return response
        else:
            return str(response)
            
    except Exception as e:
        return f"Error calling model: {str(e)}"

def create_interface():
    """Create Gradio interface"""
    
    with gr.Blocks(title="GeoLLM Model Interface", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🚀 GeoLLM Geological Intelligence Platform")
        gr.Markdown("Professional geological text analysis tool integrating triple extraction and intelligent Q&A functions")
        
        # Add tabs
        with gr.Tabs():
            # Triple extraction module
            with gr.TabItem("🔗 Triple Extraction", elem_id="triple_extraction"):
                with gr.Row():
                    with gr.Column(scale=1):
                        # Model selection area
                        gr.Markdown("## 📋 Model Configuration")
                        model_series = gr.Dropdown(
                            choices=get_model_options(),
                            value="gpt",
                            label="Model Series",
                            info="Select the model series to use"
                        )
                        
                        model_name = gr.Dropdown(
                            choices=get_common_model_names("gpt"),
                            value="gpt-3.5-turbo",
                            label="Model Name",
                            info="Select specific model name, or input manually",
                            allow_custom_value=True
                        )
                        
                        # Custom model name input box
                        custom_model_name = gr.Textbox(
                            label="Custom Model Name (Optional)",
                            placeholder="If your desired model is not in the options above, please input here",
                            info="Input here will override the selection above"
                        )
                        
                        # Prompt template selection
                        gr.Markdown("## 📝 Prompt Template")
                        prompt_template = gr.Dropdown(
                            choices=list(get_prompt_templates().keys()),
                            value="Custom",
                            label="Select Prompt Template",
                            info="Select predefined prompt template or customize"
                        )
                        
                        # Context type selection
                        gr.Markdown("## 🎯 Context Configuration")
                        context_type = gr.Dropdown(
                            choices=["No Context", "Random Context", "Best Match Context"],
                            value="No Context",
                            label="Context Type",
                            info="Choose whether to use context examples"
                        )
                        
                        num_examples = gr.Slider(
                            minimum=1,
                            maximum=3,
                            value=2,
                            step=1,
                            label="Number of Examples",
                            info="Select the number of context examples (1-3)"
                        )
                        
                    with gr.Column(scale=2):
                        # Prompt content area
                        gr.Markdown("## 🎯 Prompt Content")
                        prompt_content = gr.Textbox(
                            label="Prompt Content",
                            placeholder="Select template or customize your prompt...",
                            lines=15,
                            max_lines=25,
                            info="Will be sent to the model as system prompt"
                        )
                        
                        # User input area
                        gr.Markdown("## 💬 Geological Text Input")
                        user_content = gr.Textbox(
                            label="Geological Text to Process",
                            placeholder="Please input geological description text for triple extraction...",
                            lines=6,
                            max_lines=10
                        )
                
                # Button and output area
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                    submit_btn = gr.Button("🚀 Extract Triples", variant="primary")
                
                # Output area
                gr.Markdown("## 📤 Extraction Results")
                output = gr.Textbox(
                    label="Triple Extraction Results",
                    lines=12,
                    max_lines=20,
                    interactive=False
                )
                
                # Example area
                gr.Markdown("## 💡 Usage Examples")
                gr.Examples(
                    examples=[
                        ["gpt", "gpt-3.5-turbo", "No Context", 2, "The Noriba Gari Bao Formation originally refers to gray-green thick-bedded medium- to fine-grained lithic feldspar sandstone, feldspar quartz sandstone, feldspar sandstone occasionally interbedded with siltstone, clay rock and micritic limestone, only bivalve fossils are seen, and continuous deposition with the overlying Ninety Road Class Formation."],
                        ["gemini", "gemini-1.5-pro-002", "Random Context", 3, "The Quemo Cuo Formation has only a small outcrop in the Sewang Yongqu area in the southwest corner of the map sheet within the survey area, with an area of less than 10m2 and a thickness greater than 29.25m."],
                        ["claude", "claude-3-5-haiku-20241022", "Best Match Context", 2, "Hecosmilia sp. scabbard coral was collected from limestone; Complexastraea sp. and Radulopccten sp. scraping sea fan; Oscillopha sp., dated to the Middle Jurassic."],
                        ["deepSeek", "deepseek-ai/DeepSeek-V3", "Best Match Context", 3, "Late Triassic granite is mainly distributed in the Ladi Gongma Mianche Ri Ahri Qu area of the survey area. Regionally controlled by NW-SE trending regional faults within the structural melange zone, it is distributed in long strips. The intrusive bodies have good gregariousness and excellent zonal extensibility, with 8 exposed intrusive bodies covering an area of about 227m2."],
                    ],
                    inputs=[model_series, model_name, context_type, num_examples, user_content]
                )
                
                # Event handling
                def submit_request(series, name, custom_name, template, prompt, content, ctx_type, num_ex):
                    # Use custom model name if provided
                    final_model_name = custom_name.strip() if custom_name.strip() else name
                    return call_llm_model(series, final_model_name, prompt, content, ctx_type, num_ex)
                
                # Update model name options
                model_series.change(
                    fn=update_model_names,
                    inputs=[model_series],
                    outputs=[model_name]
                )
                
                # Update prompt content
                prompt_template.change(
                    fn=update_prompt_content,
                    inputs=[prompt_template],
                    outputs=[prompt_content]
                )
                
                # Submit button event
                submit_btn.click(
                    fn=submit_request,
                    inputs=[model_series, model_name, custom_model_name, prompt_template, prompt_content, user_content, context_type, num_examples],
                    outputs=[output]
                )
                
                # Clear button event
                clear_btn.click(
                    fn=lambda: ("", ""),
                    outputs=[user_content, output]
                )
                
                # Enter key submission
                user_content.submit(
                    fn=submit_request,
                    inputs=[model_series, model_name, custom_model_name, prompt_template, prompt_content, user_content, context_type, num_examples],
                    outputs=[output]
                )
            
            # QA module
            with gr.TabItem("❓ Intelligent Q&A", elem_id="qa_module"):
                with gr.Row():
                    with gr.Column(scale=1):
                        # Model selection area
                        gr.Markdown("## 📋 Model Configuration")
                        qa_model_series = gr.Dropdown(
                            choices=get_model_options(),
                            value="gpt",
                            label="Model Series",
                            info="Select the model series to use"
                        )
                        
                        qa_model_name = gr.Dropdown(
                            choices=get_common_model_names("gpt"),
                            value="gpt-3.5-turbo",
                            label="Model Name",
                            info="Select specific model name, or input manually",
                            allow_custom_value=True
                        )
                        
                        # Custom model name input box
                        qa_custom_model_name = gr.Textbox(
                            label="Custom Model Name (Optional)",
                            placeholder="If your desired model is not in the options above, please input here",
                            info="Input here will override the selection above"
                        )
                        
                        # QA type selection
                        gr.Markdown("## 🎯 Q&A Type")
                        qa_type = gr.Dropdown(
                            choices=["True/False", "Q&A"],
                            value="True/False",
                            label="Task Type",
                            info="Choose between judging true/false or answering questions"
                        )
                        
                        # QA Prompt template selection
                        gr.Markdown("## 📝 Prompt Template")
                        qa_prompt_template = gr.Dropdown(
                            choices=list(get_qa_prompt_templates().keys()),
                            value="Zero-shot True/False",
                            label="Select QA Prompt Template",
                            info="Select predefined prompt template or customize"
                        )
                        
                    with gr.Column(scale=2):
                        # QA Prompt content area
                        gr.Markdown("## 🎯 Prompt Content")
                        qa_prompt_content = gr.Textbox(
                            label="QA Prompt Content",
                            value="Please judge true or false based on the given text.",
                            placeholder="Select template or customize your prompt...",
                            lines=3,
                            max_lines=10,
                            info="Will be sent to the model as system prompt"
                        )
                        
                        # Geological text input area
                        gr.Markdown("## 📄 Geological Text")
                        geological_text = gr.Textbox(
                            label="Geological Background Text",
                            placeholder="Please input geological description text as background...",
                            lines=8,
                            max_lines=15,
                            info="Provides contextual information for answering questions or judging facts"
                        )
                        
                        # Question or statement input area
                        gr.Markdown("## ❓ Question/Statement")
                        question_or_statement = gr.Textbox(
                            label="Question or Statement",
                            placeholder="Please input question to answer or statement to judge...",
                            lines=3,
                            max_lines=8,
                            info="Input corresponding content based on task type"
                        )
                
                # Button and output area
                with gr.Row():
                    qa_clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                    qa_submit_btn = gr.Button("🤖 Start Q&A", variant="primary")
                
                # Output area
                gr.Markdown("## 📤 Q&A Results")
                qa_output = gr.Textbox(
                    label="Model Response",
                    lines=10,
                    max_lines=20,
                    interactive=False
                )
                
                # Example area
                gr.Markdown("## 💡 Usage Examples")
                
                # True/False examples
                with gr.Accordion("True/False Examples", open=False):
                    gr.Examples(
                        examples=[
                            ["gpt", "gpt-3.5-turbo", "True/False", "Sudden geological disasters in Huoshan County are mainly collapses, landslides, and debris flows. A total of 190 sudden geological disaster points (including hidden danger points) have been identified, including 74 collapses, 96 landslides, 14 debris flows, and 6 unstable slopes. There are 58 newly discovered geological disaster points, accounting for 30.5% of the total. Among the 190 collapses, landslides, debris flows and other sudden geological disasters in Huoshan County, most are caused by human factors. There are 163 geological disasters caused by human factors, accounting for 85.8%; there are 27 disasters formed by natural factors, accounting for 14.2%.", "In the sudden geological disasters in Huoshan County, the number of landslides exceeds the number of collapses."],
                            ["deepSeek", "deepseek-ai/DeepSeek-V3", "True/False", "Sudden geological disasters in Huoshan County are mainly collapses, landslides, and debris flows. A total of 190 sudden geological disaster points (including hidden danger points) have been identified, including 74 collapses, 96 landslides, 14 debris flows, and 6 unstable slopes.", "The total number of geological disaster points in Huoshan County exceeds 200."],
                        ],
                        inputs=[qa_model_series, qa_model_name, qa_type, geological_text, question_or_statement]
                    )
                
                # Q&A examples
                with gr.Accordion("Q&A Examples", open=False):
                    gr.Examples(
                        examples=[
                            ["gpt", "gpt-3.5-turbo", "Q&A", "Sudden geological disasters in Huoshan County are mainly collapses, landslides, and debris flows. A total of 190 sudden geological disaster points (including hidden danger points) have been identified, including 74 collapses, 96 landslides, 14 debris flows, and 6 unstable slopes. There are 58 newly discovered geological disaster points, accounting for 30.5% of the total.", "How many sudden geological disaster points are there in Huoshan County in total?"],
                            ["claude", "claude-3-5-haiku-20241022", "Q&A", "Sudden geological disasters in Huoshan County are mainly collapses, landslides, and debris flows. A total of 190 sudden geological disaster points (including hidden danger points) have been identified, including 74 collapses, 96 landslides, 14 debris flows, and 6 unstable slopes.", "Among the geological disasters in Huoshan County, which type of disaster has the largest number?"],
                        ],
                        inputs=[qa_model_series, qa_model_name, qa_type, geological_text, question_or_statement]
                    )
                
                # QA event handling
                def submit_qa_request(series, name, custom_name, q_type, template, prompt, geo_text, question):
                    # Use custom model name if provided
                    final_model_name = custom_name.strip() if custom_name.strip() else name
                    return call_qa_model(series, final_model_name, prompt, geo_text, question, q_type)
                
                def update_qa_prompt_on_type_change(qa_type_value):
                    """Update prompt template options and content when QA type changes"""
                    if qa_type_value == "True/False":
                        new_choices = ["Custom", "Zero-shot True/False", "COT True/False"]
                        new_value = "Zero-shot True/False"
                        new_prompt = "Please judge true or false based on the given text."
                        new_placeholder = "Please input statement to judge..."
                        new_label = "Statement"
                    else:  # Q&A
                        new_choices = ["Custom", "Zero-shot Q&A", "COT Q&A"]
                        new_value = "Zero-shot Q&A"
                        new_prompt = "Please answer the question based on the given text."
                        new_placeholder = "Please input question to answer..."
                        new_label = "Question"
                    
                    return (
                        gr.Dropdown(choices=new_choices, value=new_value, label="Select QA Prompt Template"),
                        gr.Textbox(value=new_prompt, label="QA Prompt Content", lines=3, max_lines=10),
                        gr.Textbox(label=new_label, placeholder=new_placeholder, lines=3, max_lines=8)
                    )
                
                # Update QA model name options
                qa_model_series.change(
                    fn=update_model_names,
                    inputs=[qa_model_series],
                    outputs=[qa_model_name]
                )
                
                # Update QA prompt content
                qa_prompt_template.change(
                    fn=update_qa_prompt_content,
                    inputs=[qa_prompt_template],
                    outputs=[qa_prompt_content]
                )
                
                # Update related components when QA type changes
                qa_type.change(
                    fn=update_qa_prompt_on_type_change,
                    inputs=[qa_type],
                    outputs=[qa_prompt_template, qa_prompt_content, question_or_statement]
                )
                
                # QA submit button event
                qa_submit_btn.click(
                    fn=submit_qa_request,
                    inputs=[qa_model_series, qa_model_name, qa_custom_model_name, qa_type, qa_prompt_template, qa_prompt_content, geological_text, question_or_statement],
                    outputs=[qa_output]
                )
                
                # QA clear button event
                qa_clear_btn.click(
                    fn=lambda: ("", "", ""),
                    outputs=[geological_text, question_or_statement, qa_output]
                )
                
                # QA enter key submission
                question_or_statement.submit(
                    fn=submit_qa_request,
                    inputs=[qa_model_series, qa_model_name, qa_custom_model_name, qa_type, qa_prompt_template, qa_prompt_content, geological_text, question_or_statement],
                    outputs=[qa_output]
                )
    
    return demo

if __name__ == "__main__":
    # Launch interface
    demo = create_interface()
    demo.launch(
        server_port=7860,         # Port number
        share=True,              # Whether to create public link
        debug=True                # Debug mode
    )