Zeggai/Question_Answer_Dataset

1from fastapi import FastAPI, HTTPException2from fastapi.responses import FileResponse3import json4from dotenv import load_dotenv5import time6import uuid7from typing import List, Dict, Optional8from datetime import datetime9from huggingface_hub import HfApi  # For file persistence in Spaces10import os11import threading12import glob13import random14from langchain_google_genai import GoogleGenerativeAI15 16# Load environment variables from .env file17load_dotenv()18 19app = FastAPI()20 21# Global variables to track generation status22generation_status = {23    "is_running": False,24    "start_time": None,25    "processed_chunks": 0,26    "total_chunks": 0,27    "questions_generated": 0,28    "completed": False,29    "result_file": None,30    "progress_file": None,  # New: track progress file31    "error": None,32    "current_api_key_index": 0,  # New: track current API key33    "failed_chunks": [],  # New: track failed chunks for retry34    "partial_results": []  # New: store partial results35}36 37generation_lock = threading.Lock()38 39def get_api_keys() -> List[str]:40    """41    Get all available Google API keys from environment variables.42    Supports GOOGLE_API_KEY, GOOGLE_API_KEY_1, GOOGLE_API_KEY_2, etc.43    """44    api_keys = []45    46    # Check for primary key47    primary_key = os.getenv("GOOGLE_API_KEY")48    if primary_key:49        api_keys.append(primary_key)50    51    # Check for numbered keys52    i = 153    while True:54        key = os.getenv(f"GOOGLE_API_KEY_{i}")55        if key:56            api_keys.append(key)57            i += 158        else:59            break60    61    if not api_keys:62        raise ValueError("No Google API keys found in environment variables")63    64    return api_keys65 66def get_next_api_key() -> tuple[str, int]:67    """68    Get the next API key in rotation and update the current index.69    Returns tuple of (api_key, key_index)70    """71    global generation_status72    73    api_keys = get_api_keys()74    75    with generation_lock:76        current_index = generation_status["current_api_key_index"]77        next_index = (current_index + 1) % len(api_keys)78        generation_status["current_api_key_index"] = next_index79    80    return api_keys[next_index], next_index81 82def save_progress_file():83    """84    Save current progress to a file that can be downloaded at any time.85    """86    global generation_status87    88    with generation_lock:89        progress_data = {90            "generation_info": {91                "status": "in_progress" if generation_status["is_running"] else "completed",92                "start_time": generation_status["start_time"],93                "processed_chunks": generation_status["processed_chunks"],94                "total_chunks": generation_status["total_chunks"],95                "questions_generated": generation_status["questions_generated"],96                "completed": generation_status["completed"],97                "current_time": datetime.utcnow().isoformat(),98                "failed_chunks": generation_status["failed_chunks"].copy(),99                "error": generation_status["error"]100            },101            "partial_dataset": {102                "dataset_info": {103                    "title": "Vaccine Guide Question-Answer Dataset (Partial)",104                    "description": "Partial dataset of question-answer pairs generated from a vaccine guide.",105                    "version": "1.1.0",106                    "created_date": generation_status["start_time"],107                    "source": "Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.pdf",108                    "generated_by": "Gemini API",109                    "total_questions": len(generation_status["partial_results"]),110                    "intended_use": "Fine-tuning medical language models for knowledge recall and reasoning",111                    "note": "This is a partial dataset. Generation may still be in progress."112                },113                "questions": generation_status["partial_results"].copy()114            }115        }116        117        # Save progress file118        progress_filename = f"vaccine_questions_progress_{int(time.time())}.json"119        generation_status["progress_file"] = progress_filename120    121    try:122        with open(f"./{progress_filename}", 'w', encoding='utf-8') as f:123            json.dump(progress_data, f, indent=4, ensure_ascii=False)124        print(f"Progress saved to {progress_filename}")125    except Exception as e:126        print(f"Error saving progress file: {e}")127 128def estimate_difficulty(question: str, q_type: str) -> str:129    """130    Estimate question difficulty based on type and content.131    132    Args:133        question (str): The question text.134        q_type (str): Question type (factual, conceptual, applied).135    136    Returns:137        str: Difficulty level (easy, medium, hard).138    """139    if q_type == "factual":140        return "easy"141    elif q_type == "conceptual":142        return "medium"143    return "hard"  # applied144 145def generate_questions_for_chunk(chunk: str, chunk_id: int, model="gemini-2.0-flash", max_retries=3) -> List[Dict]:146    """147    Generate French questions for a given document chunk using the Gemini API.148    Now includes retry logic with different API keys.149    """150    prompt = f"""151    À partir du texte suivant d'un guide sur les vaccins en français, générez 3 questions variées (factual, conceptual, applied) qui couvrent le contenu de manière exhaustive. 152    Fournissez uniquement les questions, sans réponses, en français. Retournez le résultat au format JSON, entouré de ```json\n...\n```.153    154    Texte : {chunk}155    156    Exemple de sortie :157    ```json158    [159        {{160            "question": "Combien de structures sanitaires de proximité sont impliquées dans le suivi de la vaccination ?",161            "type": "factual"162        }},163        {{164            "question": "Quel est l'impact de la réglementation de la vaccination sur la couverture vaccinale ?",165            "type": "conceptual"166        }},167        {{168            "question": "Quelles seraient les conséquences si les établissements privés ne suivaient plus la réglementation vaccinale ?",169            "type": "applied"170        }}171    ]172    ```173    """174    175    last_error = None176    177    for attempt in range(max_retries):178        try:179            # Get next API key for this attempt180            api_key, key_index = get_next_api_key()181            print(f"Chunk {chunk_id}, attempt {attempt + 1}: Using API key index {key_index}")182            183            llm = GoogleGenerativeAI(184                model=model,185                google_api_key=api_key186            )187 188            response = llm.invoke(prompt)189            190            questions_text = str(response)  # Convert response to string191            192            # Strip Markdown code fences193            if questions_text.startswith("```json\n") and questions_text.endswith("\n```"):194                questions_text = questions_text[7:-4].strip()195            elif questions_text.startswith("```") and questions_text.endswith("```"):196                questions_text = questions_text[3:-3].strip()197            198            if not questions_text:199                raise ValueError(f"Empty response for chunk {chunk_id}")200            201            questions = json.loads(questions_text)202 203            formatted_questions = []204            for q in questions:205                question_id = str(uuid.uuid4())206                difficulty = estimate_difficulty(q["question"], q["type"])207                formatted_questions.append({208                    "question_id": question_id,209                    "chunk_id": chunk_id,210                    "chunk_text": chunk,211                    "question": q["question"],212                    "type": q["type"],213                    "difficulty": difficulty,214                    "training_purpose": "Knowledge Recall" if q["type"] == "factual" else "Reasoning",215                    "validated": False,216                    "api_key_used": key_index,  # Track which key was used217                    "generation_attempt": attempt + 1218                })219            220            # Update the global status and add to partial results221            with generation_lock:222                generation_status["questions_generated"] += len(formatted_questions)223                generation_status["partial_results"].extend(formatted_questions)224            225            # Save progress after each successful chunk226            save_progress_file()227            228            print(f"Successfully generated {len(formatted_questions)} questions for chunk {chunk_id}")229            return formatted_questions230            231        except Exception as e:232            last_error = e233            print(f"Attempt {attempt + 1} failed for chunk {chunk_id}: {e}")234            235            # If this is not the last attempt, wait before retrying236            if attempt < max_retries - 1:237                wait_time = (attempt + 1) * 5  # Increasing wait time238                print(f"Waiting {wait_time} seconds before retry...")239                time.sleep(wait_time)240            241            continue242    243    # All attempts failed244    print(f"All {max_retries} attempts failed for chunk {chunk_id}. Last error: {last_error}")245    246    # Add to failed chunks list247    with generation_lock:248        generation_status["failed_chunks"].append({249            "chunk_id": chunk_id,250            "error": str(last_error),251            "attempts": max_retries252        })253    254    return []255 256def generate_questions_in_background(chunks: List[str]):257    """258    Generate questions in a background thread and update status.259    Enhanced with better error handling and progress tracking.260    """261    global generation_status262    263    try:264        all_questions = []265        266        with generation_lock:267            generation_status["total_chunks"] = len(chunks)268            generation_status["processed_chunks"] = 0269            generation_status["questions_generated"] = 0270            generation_status["partial_results"] = []271            generation_status["failed_chunks"] = []272        273        # Save initial progress file274        save_progress_file()275        276        for i, chunk in enumerate(chunks):277            print(f"Processing chunk {i+1}/{len(chunks)}...")278            questions = generate_questions_for_chunk(chunk, i)279            280            if questions:  # Only add if generation was successful281                all_questions.extend(questions)282            283            with generation_lock:284                generation_status["processed_chunks"] = i + 1285            286            # Rate limiting - slightly randomized to avoid hitting limits287            sleep_time = random.uniform(8, 11)  # Random between 8-11 seconds288            time.sleep(sleep_time)289        290        # Create final dataset291        dataset = {292            "dataset_info": {293                "title": "Vaccine Guide Question-Answer Dataset",294                "description": "A dataset of question-answer pairs generated from a vaccine guide for AI language model training.",295                "version": "1.1.0",296                "created_date": datetime.utcnow().isoformat(),297                "source": "Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.pdf",298                "generated_by": "Gemini API",299                "total_questions": len(all_questions),300                "intended_use": "Fine-tuning medical language models for knowledge recall and reasoning",301                "total_chunks_processed": len(chunks),302                "successful_chunks": len(chunks) - len(generation_status["failed_chunks"]),303                "failed_chunks": len(generation_status["failed_chunks"]),304                "failed_chunk_details": generation_status["failed_chunks"].copy()305            },306            "questions": all_questions307        }308        309        # Save the final dataset310        filename = f"vaccine_questions_final_{int(time.time())}.json"311        with open(f"./{filename}", 'w', encoding='utf-8') as f:312            json.dump(dataset, f, indent=4, ensure_ascii=False)313        314        # Update status to completed315        with generation_lock:316            generation_status["completed"] = True317            generation_status["is_running"] = False318            generation_status["result_file"] = filename319        320        # Save final progress file321        save_progress_file()322        323        success_rate = (len(chunks) - len(generation_status["failed_chunks"])) / len(chunks) * 100324        print(f"Generation completed! Success rate: {success_rate:.1f}% ({len(all_questions)} questions generated)")325    326    except Exception as e:327        print(f"Error in background generation: {e}")328        with generation_lock:329            generation_status["error"] = str(e)330            generation_status["is_running"] = False331        332        # Save progress even if there was an error333        save_progress_file()334 335def save_dataset_to_space(dataset: Dict, filename: str):336    """337    Save dataset to a file in the Space's persistent storage338    """339    persistent_path = f"./{filename}"340    with open(persistent_path, 'w', encoding='utf-8') as f:341        json.dump(dataset, f, indent=4, ensure_ascii=False)342    print(f"Dataset saved to {persistent_path}")343 344@app.get("/generate-questions")345async def generate_questions():346    """347    Endpoint to generate questions from all JSON files in the data folder348    Enhanced with multi-key support validation349    """350    global generation_status351    352    # Check if generation is already running353    with generation_lock:354        if generation_status["is_running"]:355            return {356                "status": "running",357                "message": "Generation already in progress",358                "current_status": generation_status359            }360    361    try:362        # Validate API keys before starting363        api_keys = get_api_keys()364        print(f"Found {len(api_keys)} API keys for rotation")365        366        # Reset status367        with generation_lock:368            generation_status["is_running"] = True369            generation_status["start_time"] = datetime.utcnow().isoformat()370            generation_status["processed_chunks"] = 0371            generation_status["questions_generated"] = 0372            generation_status["completed"] = False373            generation_status["result_file"] = None374            generation_status["progress_file"] = None375            generation_status["error"] = None376            generation_status["current_api_key_index"] = 0377            generation_status["failed_chunks"] = []378            generation_status["partial_results"] = []379        380        # Load all JSON files from data folder381        json_files = glob.glob("./chunk/*.json")382        383        if not json_files:384            raise HTTPException(status_code=404, detail="No JSON files found in chunk folder")385        386        all_chunks = []387        for json_file in json_files:388            with open(json_file, "r", encoding="utf-8") as f:389                chunks_data = json.load(f)390                if isinstance(chunks_data, list):391                    # If it's a list of chunks392                    for chunk in chunks_data:393                        if isinstance(chunk, dict) and "text" in chunk:394                            all_chunks.append(chunk["text"])395                        elif isinstance(chunk, str):396                            all_chunks.append(chunk)397                elif isinstance(chunks_data, dict):398                    # If it's a dict, try to extract text content399                    if "text" in chunks_data:400                        all_chunks.append(chunks_data["text"])401                    elif "content" in chunks_data:402                        all_chunks.append(chunks_data["content"])403        404        if not all_chunks:405            raise HTTPException(status_code=404, detail="No text content found in JSON files")406        407        # Start generation in background thread408        thread = threading.Thread(target=generate_questions_in_background, args=(all_chunks,))409        thread.daemon = True410        thread.start()411        412        return {413            "status": "started",414            "message": f"Question generation started for {len(json_files)} JSON files with {len(all_chunks)} chunks",415            "api_keys_available": len(api_keys),416            "current_status": generation_status417        }418    except Exception as e:419        with generation_lock:420            generation_status["is_running"] = False421            generation_status["error"] = str(e)422        raise HTTPException(status_code=500, detail=str(e))423 424@app.get("/generation-status")425async def get_generation_status():426    """427    Endpoint to check the current status of generation428    Enhanced with more detailed status information429    """430    with generation_lock:431        status_copy = generation_status.copy()432    433    # Calculate additional metrics434    if status_copy["total_chunks"] > 0:435        progress_percentage = (status_copy["processed_chunks"] / status_copy["total_chunks"]) * 100436        status_copy["progress_percentage"] = round(progress_percentage, 2)437    else:438        status_copy["progress_percentage"] = 0439    440    # Add estimated time remaining if generation is running441    if status_copy["is_running"] and status_copy["start_time"] and status_copy["processed_chunks"] > 0:442        start_time = datetime.fromisoformat(status_copy["start_time"].replace('Z', '+00:00'))443        elapsed_time = (datetime.utcnow() - start_time.replace(tzinfo=None)).total_seconds()444        chunks_per_second = status_copy["processed_chunks"] / elapsed_time if elapsed_time > 0 else 0445        446        if chunks_per_second > 0:447            remaining_chunks = status_copy["total_chunks"] - status_copy["processed_chunks"]448            estimated_remaining_seconds = remaining_chunks / chunks_per_second449            status_copy["estimated_remaining_minutes"] = round(estimated_remaining_seconds / 60, 2)450        else:451            status_copy["estimated_remaining_minutes"] = None452    453    return status_copy454 455@app.get("/download-progress")456async def download_progress():457    """458    New endpoint to download current progress at any time459    """460    global generation_status461    462    # Force save current progress463    save_progress_file()464    465    with generation_lock:466        progress_file = generation_status["progress_file"]467    468    if progress_file and os.path.exists(f"./{progress_file}"):469        return FileResponse(f"./{progress_file}", media_type="application/json", filename=progress_file)470    else:471        raise HTTPException(status_code=404, detail="No progress file available")472 473@app.get("/download/{filename}")474async def download_file(filename: str):475    """476    Endpoint to download generated files477    Enhanced with better error handling478    """479    file_path = f"./{filename}"480    if os.path.exists(file_path):481        return FileResponse(file_path, media_type="application/json", filename=filename)482    raise HTTPException(status_code=404, detail=f"File {filename} not found")483 484@app.get("/retry-failed")485async def retry_failed_chunks():486    """487    New endpoint to retry only the failed chunks488    """489    global generation_status490    491    with generation_lock:492        if generation_status["is_running"]:493            return {494                "status": "error",495                "message": "Cannot retry while generation is running"496            }497        498        failed_chunks = generation_status["failed_chunks"].copy()499    500    if not failed_chunks:501        return {502            "status": "success",503            "message": "No failed chunks to retry"504        }505    506    # This would require implementing the retry logic507    # For now, just return the failed chunks info508    return {509        "status": "info",510        "message": f"Found {len(failed_chunks)} failed chunks",511        "failed_chunks": failed_chunks,512        "note": "Retry functionality can be implemented based on requirements"513    }514 515@app.get("/api-keys-status")516async def get_api_keys_status():517    """518    New endpoint to check API keys status519    """520    try:521        api_keys = get_api_keys()522        return {523            "status": "success",524            "total_keys": len(api_keys),525            "current_key_index": generation_status["current_api_key_index"],526            "message": f"{len(api_keys)} API keys configured for rotation"527        }528    except Exception as e:529        return {530            "status": "error",531            "message": str(e)532        }533 534@app.get("/")535async def root():536    """537    Root endpoint that serves the HTML UI from the index.html file.538    """539    print("Serving index.html") # Debug log to confirm serving540    return FileResponse("./index.html", media_type="text/html")541 542if __name__ == "__main__":543    import uvicorn544    uvicorn.run(app, host="0.0.0.0", port=7860)