david167/question-generation-api

1import os2import logging3import torch4from transformers import AutoTokenizer, AutoModelForCausalLM5import gradio as gr6import json7import re8 9# Configure logging10logging.basicConfig(level=logging.INFO)11logger = logging.getLogger(__name__)12 13class ModelManager:14    def __init__(self):15        self.model = None16        self.tokenizer = None17        self.device = None18        self.model_loaded = False19        self.load_model()20 21    def load_model(self):22        """Load the model and tokenizer"""23        try:24            logger.info("Starting model loading...")25            26            # Check if CUDA is available27            if torch.cuda.is_available():28                torch.cuda.set_device(0)29                self.device = "cuda:0"30            else:31                self.device = "cpu"32            logger.info(f"Using device: {self.device}")33            34            if self.device == "cuda:0":35                logger.info(f"GPU: {torch.cuda.get_device_name()}")36                logger.info(f"VRAM Available: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")37            38            # Get HF token from environment39            hf_token = os.getenv("HF_TOKEN")40            41            logger.info("Loading Llama-3.1-8B-Instruct model...")42            base_model_name = "meta-llama/Llama-3.1-8B-Instruct"43            44            self.tokenizer = AutoTokenizer.from_pretrained(45                base_model_name,46                use_fast=True,47                trust_remote_code=True,48                token=hf_token49            )50            51            self.model = AutoModelForCausalLM.from_pretrained(52                base_model_name,53                torch_dtype=torch.float16 if self.device == "cuda:0" else torch.float32,54                device_map="auto" if self.device == "cuda:0" else None,55                trust_remote_code=True,56                token=hf_token57            )58            59            # Set pad token60            if self.tokenizer.pad_token is None:61                self.tokenizer.pad_token = self.tokenizer.eos_token62            63            self.model_loaded = True64            logger.info("✅ Model loaded successfully!")65            66        except Exception as e:67            logger.error(f"❌ Error loading model: {str(e)}")68            self.model_loaded = False69 70def generate_response(prompt, temperature=0.8, model_manager=None):71    """ELEGANT AI ARCHITECT SOLUTION - Clean, simple, effective"""72    if not model_manager or not model_manager.model_loaded:73        return "Model not loaded"74 75    try:76        # Detect request type77        is_cot_request = any(phrase in prompt.lower() for phrase in [78            "return exactly this json array",79            "chain of thinking", 80            "verbatim",81            "json array (no other text)"82        ])83        84        # Get actual model context85        max_context = getattr(model_manager.model.config, "max_position_embeddings", 8192)86        logger.info(f"Model context: {max_context} tokens")87        88        # SIMPLE, CLEAR PROMPT FORMATTING89        if is_cot_request:90            system_msg = "You are an expert at generating JSON training data. Return only valid JSON arrays as requested, no additional text."91        else:92            system_msg = "You are a helpful AI assistant generating high-quality training data."93            94        formatted_prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>95 96{system_msg}97 98<|eot_id|><|start_header_id|>user<|end_header_id|>99 100{prompt}101 102<|eot_id|><|start_header_id|>assistant<|end_header_id|>103 104"""105        106        # SMART TOKEN ALLOCATION107        if is_cot_request:108            # CoT needs substantial output for complete JSON109            max_new_tokens = 3000  # Generous but not excessive  110            min_new_tokens = 500   # Ensure JSON completion111        else:112            max_new_tokens = 1500113            min_new_tokens = 50114            115        # Reserve space for input116        max_input_tokens = max_context - max_new_tokens - 100117        118        logger.info(f"Token plan: Input≤{max_input_tokens}, Output={min_new_tokens}-{max_new_tokens}")119 120        # Tokenize121        inputs = model_manager.tokenizer(122            formatted_prompt,123            return_tensors="pt",124            truncation=True,125            max_length=max_input_tokens126        )127        128        # Move to device129        if model_manager.device == "cuda:0":130            inputs = {k: v.to(next(model_manager.model.parameters()).device) for k, v in inputs.items()}131        132        # CLEAN GENERATION133        with torch.no_grad():134            outputs = model_manager.model.generate(135                **inputs,136                max_new_tokens=max_new_tokens,137                min_new_tokens=min_new_tokens,138                temperature=temperature,139                top_p=0.9,140                do_sample=True,141                pad_token_id=model_manager.tokenizer.eos_token_id,142                early_stopping=False,143                repetition_penalty=1.1144            )145        146        # Decode147        full_response = model_manager.tokenizer.decode(outputs[0], skip_special_tokens=True)148        149        # Log stats150        input_len = inputs['input_ids'].shape[1]151        output_len = outputs[0].shape[0]152        generated_len = output_len - input_len153        logger.info(f"Generated {generated_len} tokens (min was {min_new_tokens})")154        155        # CLEAN EXTRACTION156        if "<|start_header_id|>assistant<|end_header_id|>" in full_response:157            response = full_response.split("<|start_header_id|>assistant<|end_header_id|>", 1)[-1].strip()158        else:159            # Fallback160            response = full_response[len(formatted_prompt):].strip()161        162        # For CoT, extract clean JSON if possible163        if is_cot_request and '[' in response and ']' in response:164            # Find the most complete JSON array165            json_pattern = r'\[(?:[^[\]]+|\[[^\]]*\])*\]'166            matches = re.findall(json_pattern, response, re.DOTALL)167            168            if matches:169                # Pick the longest match (most complete)170                best_match = max(matches, key=len)171                # Verify it has reasonable content172                if '"user"' in best_match and '"assistant"' in best_match:173                    logger.info(f"Extracted JSON: {len(best_match)} chars")174                    response = best_match175        176        logger.info(f"Final response: {len(response)} chars")177        return response.strip()178 179    except Exception as e:180        logger.error(f"Generation error: {e}")181        return f"Error: {e}"182 183# Initialize model184model_manager = ModelManager()185 186def respond(message, history, temperature):187    """Gradio interface function"""188    try:189        response = generate_response(message, temperature, model_manager)190        history.append([message, response])191        return history, ""192    except Exception as e:193        logger.error(f"Error in respond: {e}")194        history.append([message, f"Error: {e}"])195        return history, ""196 197# Create Gradio interface198with gr.Blocks(title="Question Generation API") as demo:199    gr.Markdown("# Question Generation API")200    201    chatbot = gr.Chatbot(height=400)202    msg = gr.Textbox(label="Message", placeholder="Enter your prompt...")203    temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.1, label="Temperature")204    205    with gr.Row():206        submit = gr.Button("Submit", variant="primary")207        clear = gr.Button("Clear")208    209    submit.click(respond, [msg, chatbot, temperature], [chatbot, msg])210    msg.submit(respond, [msg, chatbot, temperature], [chatbot, msg])211    clear.click(lambda: ([], ""), outputs=[chatbot, msg])212 213if __name__ == "__main__":214    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)