Update main.py

main.py

@@ -2,19 +2,33 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import json
+import re
 
 # --- App and Model Loading ---
 app = FastAPI()
-
 model_name = "Qwen/Qwen2-0.5B-Instruct"
 print("Loading model...")
-# To leverage a GPU on Hugging Face Spaces, device_map="auto" is key
+
+# Optimized model loading
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype="auto",
-    device_map="auto"
+    torch_dtype=torch.float16,  # Use float16 for speed
+    device_map="auto",
+    trust_remote_code=True,
+    use_cache=True,  # Enable KV cache for faster generation
+    low_cpu_mem_usage=True
+)
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    padding_side="left"  # Better for batch generation
 )
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Set pad token if not exists
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
 print("Model loaded successfully.")
 
 # --- API Request and Response Models ---
@@ -25,53 +39,128 @@ class GenerationRequest(BaseModel):
 class GenerationResponse(BaseModel):
     data: list
 
-# --- API Endpoint ---
-@app.post("/generate", response_model=GenerationResponse)
-async def generate_data(request: GenerationRequest):
-    prompt = f"""
-You are a data generator. Your task is to generate {request.batch_size} random, non-similar rows of data based on the following commands.
-Each command corresponds to a column.
-Commands: {request.llm_commands}
-Return the data as a valid JSON array of arrays, where each inner array represents a row.
-For example, for the commands ["an age between 20 and 30", "a random city in California"], the output should look like:
-[[25, "Los Angeles"], [22, "San Francisco"]]
-Do not include any extra text, explanations, or markdown formatting in your response. Only output the raw JSON array.
-"""
-
-    messages = [
-        {"role": "system", "content": "You are a helpful assistant that generates structured data."},
-        {"role": "user", "content": prompt}
-    ]
+# --- Helper Functions ---
+def extract_json_from_text(text: str):
+    """Extract JSON array from model output, handling extra text."""
+    # Look for JSON array pattern
+    json_pattern = r'\[\s*\[.*?\]\s*\]'
+    matches = re.findall(json_pattern, text, re.DOTALL)
+    
+    if matches:
+        try:
+            return json.loads(matches[0])
+        except:
+            pass
+    
+    # Fallback: try to find anything that looks like nested arrays
+    try:
+        # Find content between first [ and last ]
+        start = text.find('[')
+        end = text.rfind(']') + 1
+        if start != -1 and end != 0:
+            json_candidate = text[start:end]
+            return json.loads(json_candidate)
+    except:
+        pass
+    
+    return None
 
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+def create_optimized_prompt(commands: list[str], batch_size: int) -> str:
+    """Create a more structured prompt to reduce hallucination."""
+    return f"""Generate exactly {batch_size} rows of data. Each row has {len(commands)} columns:
+{chr(10).join([f'Column {i+1}: {cmd}' for i, cmd in enumerate(commands)])}
 
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+Output format: JSON array only, no explanations.
+Example: [[value1, value2], [value3, value4]]
 
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=2048 # Increased to handle larger batches
-    )
-    
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
-    
-    response_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+Generate {batch_size} rows:"""
 
+# --- API Endpoint ---
+@app.post("/generate", response_model=GenerationResponse)
+async def generate_data(request: GenerationRequest):
     try:
-        # The model might still add extra text, so we clean it
-        json_response = torch.tensor(eval(response_text.strip()))
-        return {"data": json_response.tolist()}
+        # Create optimized prompt
+        prompt = create_optimized_prompt(request.llm_commands, request.batch_size)
+        
+        messages = [
+            {"role": "system", "content": "You are a precise data generator. Output only valid JSON arrays with no extra text."},
+            {"role": "user", "content": prompt}
+        ]
+        
+        # Apply chat template
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        
+        # Tokenize with optimized settings
+        model_inputs = tokenizer(
+            text, 
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048,  # Limit input length
+            padding=False
+        ).to(model.device)
+        
+        # Generate with optimized parameters
+        with torch.no_grad():  # Disable gradients for inference
+            generated_ids = model.generate(
+                **model_inputs,
+                max_new_tokens=min(1024, request.batch_size * 20),  # Dynamic max tokens
+                min_new_tokens=10,
+                do_sample=True,
+                temperature=0.7,  # Balanced creativity/consistency
+                top_p=0.9,
+                top_k=50,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                use_cache=True,
+                num_beams=1,  # Faster than beam search
+                early_stopping=True
+            )
+        
+        # Extract generated text
+        generated_ids = [
+            output_ids[len(input_ids):] 
+            for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
+        
+        response_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        print(f"Raw model output: {response_text[:200]}...")  # Debug print
+        
+        # Extract JSON data
+        json_data = extract_json_from_text(response_text)
+        
+        if json_data and isinstance(json_data, list):
+            # Validate data structure
+            if len(json_data) > 0 and isinstance(json_data[0], list):
+                # Ensure we have the right number of columns
+                expected_cols = len(request.llm_commands)
+                filtered_data = [
+                    row for row in json_data 
+                    if isinstance(row, list) and len(row) == expected_cols
+                ]
+                
+                if filtered_data:
+                    return {"data": filtered_data[:request.batch_size]}
+        
+        print(f"Failed to parse JSON. Raw output: {response_text}")
+        return {"data": []}
+        
     except Exception as e:
-        print(f"Error parsing model output: {e}")
-        print(f"Raw output was: {response_text}")
-        # Return empty on failure to prevent crashing the Inngest job
+        print(f"Error in generation: {e}")
         return {"data": []}
 
 @app.get("/")
 def read_root():
-    return {"status": "ok"}
+    return {"status": "ok", "model": model_name}
+
+@app.get("/health")
+def health_check():
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None,
+        "device": str(model.device) if model else "unknown"
+    }