Gradio_en.py

# -*- coding: utf-8 -*-
import gradio as gr
import pandas as pd
import json
import random
import numpy as np
import faiss
from transformers import AutoTokenizer, AutoModel
import torch
from collections import defaultdict
from LLM import zero_shot
from prompt_generate import generate_prompt_with_examples as generate_prompt

_knn_retriever = None

class EntityLevelRetriever:
    def __init__(self, model_name='bert-base-uncased'):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.index = faiss.IndexFlatL2(768)  
        self.entity_db = []
        self.metadata = []
        self.train_data = []

    def _get_entity_span(self, text, entity):
        start = text.find(entity)
        if start == -1:
            return None
        return (start, start + len(entity))

    def _generate_entity_embedding(self, text, entity):
        span = self._get_entity_span(text, entity)
        if not span:
            return None

        inputs = self.tokenizer(text, return_tensors='pt', truncation=True)
        with torch.no_grad():
            outputs = self.model(**inputs)

        char_to_token = lambda x: inputs.char_to_token(x)
        start_token = char_to_token(span[0])
        end_token = char_to_token(span[1]-1)

        if not start_token or not end_token:
            return None

        entity_embedding = outputs.last_hidden_state[0, start_token:end_token+1].mean(dim=0).numpy()
        return entity_embedding.astype('float32')

    def build_index(self, train_path):
        with open(train_path, 'r', encoding='utf-8') as f:
            dataset = json.load(f)

        self.train_data = dataset[500:1000]
        
        for item in self.train_data:
            text = item['text']
            for triple in item['triple_list']:
                for entity in [triple[0], triple[2]]:
                    embedding = self._generate_entity_embedding(text, entity)
                    if embedding is not None:
                        self.entity_db.append(embedding)
                        self.metadata.append({
                            'entity': entity,
                            'type': triple[1], 
                            'context': text,
                            'full_item': item
                        })

        if self.entity_db:
            self.index.add(np.array(self.entity_db))

    def search_similar_texts(self, query_text, top_k=3):
        """Search for similar texts based on entity embeddings"""
        if not self.train_data:
            return []
            
        query_entities = self._extract_potential_entities(query_text)
        
        context_scores = defaultdict(float)
        context_items = {}
        
        for entity in query_entities:
            embedding = self._generate_entity_embedding(query_text, entity)
            if embedding is None:
                continue
                
            distances, indices = self.index.search(np.array([embedding]), top_k * 2)
            
            for j in range(len(indices[0])):
                idx = indices[0][j]
                if 0 <= idx < len(self.metadata):
                    ctx_info = self.metadata[idx]
                    distance = distances[0][j]
                    context = ctx_info['context']
                    
                    # Weight by inverse distance
                    score = 1 / (1 + distance)
                    context_scores[context] += score
                    context_items[context] = ctx_info['full_item']
        
        # Sort by score and return top_k
        scored_contexts = sorted(context_scores.items(), key=lambda x: x[1], reverse=True)
        
        results = []
        for context, score in scored_contexts[:top_k]:
            if context in context_items:
                results.append(context_items[context])
        
        return results
    
    def _extract_potential_entities(self, text):
        """Simple entity extraction - you can improve this with better NER"""
        # Split by common delimiters and filter meaningful terms
        import re
        # Extract potential geological terms (capitalized words, formations, etc.)
        entities = re.findall(r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b', text)
        # Also include technical terms
        entities.extend(re.findall(r'\b\w+(?:stone|rock|formation|group|member)\b', text, re.IGNORECASE))
        return list(set(entities))

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 prompting": """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: "Lithology", "Paleontological", "Thickness", "Outcrop", "Develop", "Exposed", "Distribution pattern", "Contained", "Coordinates", "Age", "Integration of contacts", "Unconformity contact", "Fault contact", "Belongs to", "Elevation", "Located at", "Exposed area", "Geotectonic position", "Stratigraphical zoning", "Consolidated contact", "Administrative area", "Engulfed", "Length", "Invade", 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 prompting": """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: "Lithology", "Paleontological", "Thickness", "Outcrop", "Develop", "Exposed", "Distribution pattern", "Contained", "Coordinates", "Age", "Integration of contacts", "Unconformity contact", "Fault contact", "Belongs to", "Elevation", "Located at", "Exposed area", "Geotectonic position", "Stratigraphical zoning", "Consolidated contact", "Administrative area", "Engulfed", "Length", "Invade", 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:
Exposed: Refers to rocks or strata exposed at the surface or near-surface, not covered or buried. Example: (Late Ordovician-Silurian intrusive rocks, Exposed, southern investigation area).

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

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

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

Consolidated contact: Specifically refers to parallel unconformity types with consistent attitudes, emphasizing depositional sequence interruption but without structural deformation. Example: (Solake Formation, Consolidated 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).

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

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

Outcrop: Specifically refers to actually exposed stratigraphic entities in a region, emphasizing observable surface geological units. Example: (Hongliugou gold-copper mining area, Outcrop, 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).

Exposed area: Characterizes horizontal distribution range of geological units, presented in standardized form combining numerical values and units. Example: (intrusive rocks, Exposed 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.

Contained: 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, Contained, chert bands).

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

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

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

Paleontological: Records fossil biological information occurring in strata, requiring complete Latin scientific names and classification features. Example: (strata, Paleontological, 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).

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

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

Invade: Describes geological processes of magmatic rock bodies penetrating country rocks, including accompanying phenomena such as contact metamorphism. Example: (Gaozhou Shell Stone Formation, Invade, 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": "",
        "Yes/No QA": "Please judge true or false based on the given text.",
        "Factoid QA": "Please answer the question based on the given text.",
        "CoT prompting of Yes/No QA": "Please first judge true or false, and provide your reasoning basis.",
        "CoT prompting of Factoid QA": "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 initialize_knn_retriever():
    """Initialize the KNN retriever"""
    global _knn_retriever
    if _knn_retriever is None:
        try:
            # print("Initializing KNN retriever...")
            _knn_retriever = EntityLevelRetriever()
            _knn_retriever.build_index('./data/train_triples.json')
            # print("KNN retriever initialized successfully!")
        except Exception as e:
            # print(f"Failed to initialize KNN retriever: {e}")
            _knn_retriever = None
    return _knn_retriever is not None

def generate_random_context_prompt(user_text, num_examples):
    """Generate Random sampling prompts"""
    if not load_train_data():
        return "Unable to load training data"
    
    try:
        random_prompt = generate_prompt(_text_series, _label_series, num_examples)
        print(random_prompt)
        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 sampling prompt: {e}"

def generate_knn_context_prompt(user_text, num_examples):
    """Generate KNN-based sampling prompts"""
    global _knn_retriever
    
    if not initialize_knn_retriever():
        return "Unable to initialize KNN retriever"
    
    try:
        similar_items = _knn_retriever.search_similar_texts(user_text, num_examples)
        
        if not similar_items:
            return f"No similar examples found, performing zero-shot extraction:\n{user_text}"
        
        examples_text = ""
        for i, item in enumerate(similar_items):
            examples_text += f"Example {i+1}:\n"
            examples_text += f"Text: {item['text']}\n"
            examples_text += f"Triples: {json.dumps(item['triple_list'], ensure_ascii=False)}\n\n"
        print(examples_text)
        return f"Here are geological description text and triple extraction examples based on KNN similarity:\n\n{examples_text}Please extract triples based on the examples:\n{user_text}"
    
    except Exception as e:
        return f"Failed to generate KNN-based sampling 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="Select the specific model", 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 ", 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 ", 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 Demonstration type
        if context_type == "No demonstration":
            if prompt_content.strip():
                full_content = prompt_content.strip() + "\n\n" + user_content
            else:
                full_content = user_content
        elif context_type == "Random sampling":
            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 == "KNN-based sampling":
            context_prompt = generate_knn_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 == "Yes/No QA":
                return "Please input factual statement to judge"
            else:
                return "Please input question to answer"
        
        # Combine complete input content
        if qa_type == "Yes/No QA":
            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:  # Factoid QA
            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-Toolkit: An Interactive Platform for Geological Text Understanding and Knowledge Extraction")
        gr.Markdown("A domain-specific, modular framework designed to operationalize large language models (LLMs) for advanced geological natural language processing")
        
        # Add tabs
        with gr.Tabs():
            # Triple extraction module
            with gr.TabItem("🔗 KG 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="Select the specific model",
                            info="Select specific model name",
                            allow_custom_value=True
                        )
                        
                        # 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"
                        )
                        
                        # Demonstration type selection
                        gr.Markdown("## 🎯 Few-Shot Prompting Settings")
                        context_type = gr.Dropdown(
                            choices=["No demonstration", "Random sampling", "KNN-based sampling"],
                            value="No demonstration",
                            label="Demonstration type",
                            info="Choose the method for selecting 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  area
                        gr.Markdown("## 🎯 Prompt")
                        prompt_content = gr.Textbox(
                            label="Prompt ",
                            placeholder="Select template or customize your prompt...",
                            lines=15,
                            max_lines=25,
                            info="Defines the task-specific prompt during geological NLP tasks"
                        )
                        
                        # 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=8,
                    max_lines=10,
                    interactive=False
                )
                
                # Example area
                gr.Markdown("## 💡 Usage Examples")
                gr.Examples(
                    examples=[
                        ["gpt", "gpt-3.5-turbo", "No demonstration", 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 sampling", 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", "KNN-based sampling", 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", "KNN-based sampling", 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, template, prompt, content, ctx_type, num_ex):
                    # Use the selected model name directly
                    return call_llm_model(series, 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 
                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, 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, prompt_template, prompt_content, user_content, context_type, num_examples],
                    outputs=[output]
                )
            
            # QA module
            with gr.TabItem("❓ Geological 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="Select the specific model",
                            info="Select specific model name",
                            allow_custom_value=True
                        )
                        
                        # QA type selection
                        gr.Markdown("## 🎯 Geological Q&A")
                        qa_type = gr.Dropdown(
                            choices=["Yes/No QA", "Factoid QA"],
                            value="Yes/No QA",
                            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="Yes/No QA",
                            label="Select QA Prompt Template",
                            info="Select predefined prompt template or customize"
                        )
                        
                    with gr.Column(scale=2):
                        # QA Prompt  area
                        gr.Markdown("## 🎯 Prompt")
                        qa_prompt_content = gr.Textbox(
                            label="QA Prompt ",
                            value="Please judge true or false based on the given text.",
                            placeholder="Select template or customize your prompt...",
                            lines=3,
                            max_lines=10,
                            info="Defines the task-specific prompt during geological NLP tasks"
                        )
                        
                        # Geological text input area
                        gr.Markdown("## 📄 Geological Text")
                        geological_text = gr.Textbox(
                            label="Contextual Text in Geological Domain",
                            placeholder="Please input geological description text as background...",
                            lines=8,
                            max_lines=10,
                            info="Provides contextual information for Q&A tasks"
                        )
                        
                        # 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=10,
                    interactive=False
                )
                
                # Example area
                gr.Markdown("## 💡 Usage Examples")
                
                # Yes/No QA examples
                with gr.Accordion("Yes/No QA Examples", open=False):
                    gr.Examples(
                        examples=[
                            ["gpt", "gpt-3.5-turbo", "Yes/No QA", "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", "Yes/No QA", "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]
                    )
                
                # Factoid QA examples
                with gr.Accordion("Factoid QA Examples", open=False):
                    gr.Examples(
                        examples=[
                            ["gpt", "gpt-3.5-turbo", "Factoid QA", "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", "Factoid QA", "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, q_type, template, prompt, geo_text, question):

                    return call_qa_model(series, 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 == "Yes/No QA":
                        new_choices = ["Custom", "Yes/No QA", "CoT prompting of Yes/No QA"]
                        new_value = "Yes/No QA"
                        new_prompt = "Please judge true or false based on the given text."
                        new_placeholder = "Please input statement to judge..."
                        new_label = "Statement"
                    else:  # Factoid QA
                        new_choices = ["Custom", "Factoid QA", "CoT prompting of Factoid QA"]
                        new_value = "Factoid QA"
                        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 ", 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 
                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_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回车键提交
                question_or_statement.submit(
                    fn=submit_qa_request,
                    inputs=[qa_model_series, qa_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 with password protection
    demo = create_interface()
    demo.launch(
        server_port=7860,
        share=True,
        debug=True,
        # auth=("geollm", "research2025"), 
        # auth_message="Please enter credentials to access GeoLLM Geological Intelligence Platform"
    )