david167/question-generation-api
0
1import os2import logging3import torch4from transformers import AutoTokenizer, AutoModelForCausalLM5import gradio as gr6 7# Configure logging8logging.basicConfig(level=logging.INFO)9logger = logging.getLogger(__name__)10 11class ModelManager:12 def __init__(self):13 self.model = None14 self.tokenizer = None15 self.device = None16 self.model_loaded = False17 self.load_model()18 19 def load_model(self):20 """Load the model and tokenizer"""21 try:22 logger.info("Starting model loading...")23 24 # Check if CUDA is available25 if torch.cuda.is_available():26 torch.cuda.set_device(0)27 self.device = "cuda:0"28 else:29 self.device = "cpu"30 logger.info(f"Using device: {self.device}")31 32 if self.device == "cuda:0":33 logger.info(f"GPU: {torch.cuda.get_device_name()}")34 logger.info(f"VRAM Available: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")35 36 # Get HF token from environment37 hf_token = os.getenv("HF_TOKEN")38 39 logger.info("Loading Llama-3.1-8B-Instruct model...")40 base_model_name = "meta-llama/Llama-3.1-8B-Instruct"41 42 self.tokenizer = AutoTokenizer.from_pretrained(43 base_model_name,44 use_fast=True,45 trust_remote_code=True,46 token=hf_token47 )48 49 self.model = AutoModelForCausalLM.from_pretrained(50 base_model_name,51 torch_dtype=torch.float16 if self.device == "cuda:0" else torch.float32,52 device_map={"": 0} if self.device == "cuda:0" else None,53 trust_remote_code=True,54 low_cpu_mem_usage=True,55 use_safetensors=True,56 token=hf_token57 )58 59 if self.device == "cuda:0":60 self.model = self.model.to(self.device)61 62 self.model_loaded = True63 logger.info("Model loaded successfully!")64 65 except Exception as e:66 logger.error(f"Error loading model: {str(e)}")67 self.model_loaded = False68 69# Initialize model manager70model_manager = ModelManager()71 72def generate_response(prompt, temperature=0.8):73 """Simple function to generate a response from a prompt"""74 if not model_manager.model_loaded:75 return "Model not loaded yet. Please wait..."76 77 try:78 # Create the Llama-3.1 chat format79 formatted_prompt = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>80 81{prompt}82 83<|eot_id|><|start_header_id|>assistant<|end_header_id|>84 85"""86 87 # Determine context window and USE ABSOLUTE MAXIMUM88 try:89 max_ctx = getattr(model_manager.model.config, "max_position_embeddings", 131072) # Llama 3.1 supports up to 131k90 except Exception:91 max_ctx = 131072 # Use maximum possible92 93 logger.info(f"Model max context: {max_ctx} tokens")94 95 # Detect if this is a Chain of Thinking request96 is_cot_request = ("chain-of-thinking" in prompt.lower() or 97 "chain of thinking" in prompt.lower() or98 "Return exactly this JSON array" in prompt or99 ("verbatim" in prompt.lower() and "json array" in prompt.lower()))100 101 # MAXIMIZE GENERATION TOKENS - use most of context for generation102 if is_cot_request:103 # For CoT, use MAXIMUM possible generation tokens104 gen_max_new_tokens = 16384 # Very high limit for complete responses105 min_tokens = 2000 # High minimum to force complete generation106 # Allow most of context for input107 allowed_input_tokens = max_ctx - gen_max_new_tokens - 100 # Small safety buffer108 logger.info(f"CoT REQUEST - MAXIMIZED: min_tokens={min_tokens}, max_new_tokens={gen_max_new_tokens}, input_limit={allowed_input_tokens}")109 else:110 # Standard requests111 gen_max_new_tokens = 8192112 min_tokens = 200113 allowed_input_tokens = max_ctx - gen_max_new_tokens - 100114 115 # Tokenize the input with safe truncation116 inputs = model_manager.tokenizer(117 formatted_prompt,118 return_tensors="pt",119 truncation=True,120 max_length=allowed_input_tokens121 )122 123 # Move inputs to the same device as the model124 if model_manager.device == "cuda:0":125 model_device = next(model_manager.model.parameters()).device126 inputs = {k: v.to(model_device) for k, v in inputs.items()}127 128 # Generate response with MAXIMUM settings129 with torch.no_grad():130 outputs = model_manager.model.generate(131 **inputs,132 max_new_tokens=gen_max_new_tokens,133 min_new_tokens=min_tokens,134 temperature=temperature,135 top_p=0.95,136 do_sample=True,137 num_beams=1,138 pad_token_id=model_manager.tokenizer.eos_token_id,139 eos_token_id=model_manager.tokenizer.eos_token_id,140 early_stopping=False, # Never stop early141 repetition_penalty=1.05,142 no_repeat_ngram_size=0,143 length_penalty=1.0,144 # Force generation to continue145 use_cache=True146 )147 148 # Decode the response149 generated_text = model_manager.tokenizer.decode(outputs[0], skip_special_tokens=True)150 151 # Log generation details for debugging152 input_length = inputs['input_ids'].shape[1]153 output_length = outputs[0].shape[0]154 generated_length = output_length - input_length155 logger.info(f"Generation stats - Input: {input_length} tokens, Generated: {generated_length} tokens, Min required: {min_tokens}")156 157 if generated_length < min_tokens:158 logger.warning(f"Generated {generated_length} tokens but minimum was {min_tokens} - response may be truncated")159 160 # Post-decode guard: if a top-level JSON array closes, trim to the first full array161 # This helps prevent trailing prose like 'assistant' or 'Message'.162 try:163 # Track both bracket and brace depth to find first complete JSON structure164 bracket_depth = 0 # [ ]165 brace_depth = 0 # { }166 in_string = False167 escape_next = False168 start_idx = None169 end_idx = None170 171 for i, ch in enumerate(generated_text):172 # Handle string escaping173 if escape_next:174 escape_next = False175 continue176 177 if ch == '\\':178 escape_next = True179 continue180 181 # Track if we're inside a string182 if ch == '"' and not escape_next:183 in_string = not in_string184 continue185 186 # Only count brackets/braces outside of strings187 if not in_string:188 if ch == '[':189 if bracket_depth == 0 and brace_depth == 0 and start_idx is None:190 start_idx = i191 bracket_depth += 1192 elif ch == ']':193 bracket_depth = max(0, bracket_depth - 1)194 if bracket_depth == 0 and brace_depth == 0 and start_idx is not None:195 end_idx = i196 break197 elif ch == '{':198 brace_depth += 1199 elif ch == '}':200 brace_depth = max(0, brace_depth - 1)201 202 if start_idx is not None and end_idx is not None and end_idx > start_idx:203 # Extract just the complete JSON array204 json_text = generated_text[start_idx:end_idx+1]205 logger.info(f"Extracted complete JSON array of length {len(json_text)}")206 generated_text = json_text207 elif start_idx is not None:208 # Found start but no end - response was truncated209 logger.warning("JSON array started but never closed - response truncated")210 # Try to extract what we have and let the client handle it211 generated_text = generated_text[start_idx:]212 except Exception as e:213 logger.warning(f"Error in JSON extraction: {e}")214 pass215 216 # Extract just the assistant's response217 if "<|start_header_id|>assistant<|end_header_id|>" in generated_text:218 response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()219 else:220 # Better fallback: look for the start of actual content (JSON or text)221 import re222 223 # Look for JSON array or object start224 json_match = re.search(r'(\[|\{)', generated_text)225 if json_match and json_match.start() > len(formatted_prompt) // 2:226 response = generated_text[json_match.start():].strip()227 else:228 # Look for the end of the prompt pattern229 prompt_end_patterns = [230 "<|end_header_id|>",231 "<|eot_id|>",232 "assistant",233 "\n\n"234 ]235 236 response = generated_text237 for pattern in prompt_end_patterns:238 if pattern in generated_text:239 parts = generated_text.split(pattern)240 if len(parts) > 1:241 # Take the last substantial part242 candidate = parts[-1].strip()243 if len(candidate) > 20: # Ensure it's not too short244 response = candidate245 break246 247 # Ultimate fallback - just return everything after a reasonable point248 if response == generated_text:249 # Skip approximately the prompt length but be conservative250 skip_chars = min(len(formatted_prompt) // 2, len(generated_text) // 3)251 response = generated_text[skip_chars:].strip()252 253 logger.info(f"Generated response length: {len(response)} characters")254 return response255 256 except Exception as e:257 logger.error(f"Error generating response: {str(e)}")258 return f"Error: {str(e)}"259 260def respond(message, history, temperature):261 """Gradio interface function for chat"""262 response = generate_response(message, temperature)263 264 # Update history265 history.append({"role": "user", "content": message})266 history.append({"role": "assistant", "content": response})267 268 return history, ""269 270# Create the Gradio interface271with gr.Blocks(title="Question Generation API") as demo:272 gr.Markdown("# Simple LLM API")273 gr.Markdown("Send a prompt and get a response. No templates, just direct model interaction.")274 275 with gr.Row():276 with gr.Column(scale=4):277 chatbot = gr.Chatbot(278 label="Chat",279 type="messages",280 height=400281 )282 msg = gr.Textbox(283 label="Message",284 placeholder="Enter your prompt here...",285 lines=3286 )287 with gr.Row():288 submit = gr.Button("Send", variant="primary")289 clear = gr.Button("Clear")290 291 with gr.Column(scale=1):292 temperature = gr.Slider(293 minimum=0.1,294 maximum=2.0,295 value=0.8,296 step=0.1,297 label="Temperature",298 info="Higher = more creative"299 )300 gr.Markdown("""301 ### API Usage302 This model accepts any prompt and returns a response.303 304 For JSON responses, include instructions in your prompt like:305 - "Return as a JSON array"306 - "Format as JSON"307 - "List as JSON"308 309 The model will follow your instructions.310 """)311 312 # Set up event handlers313 submit.click(respond, [msg, chatbot, temperature], [chatbot, msg])314 msg.submit(respond, [msg, chatbot, temperature], [chatbot, msg])315 clear.click(lambda: ([], ""), outputs=[chatbot, msg])316 317if __name__ == "__main__":318 demo.launch(319 server_name="0.0.0.0",320 server_port=7860,321 share=False322 )