Create app.py

app.py

@@ -0,0 +1,148 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import json
+import random
+from typing import List, Optional
+
+app = FastAPI(title="Qwen Data Generator API")
+
+# Load model and tokenizer
+model_name = "Qwen/Qwen2.5-3B-Instruct"
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+print("Model loaded successfully!")
+
+class GenerationRequest(BaseModel):
+    llm_commands: List[str]
+    batch_size: int = 50
+    seed: Optional[int] = None
+
+class GenerationResponse(BaseModel):
+    success: bool
+    data: List[List[str]]
+    error: Optional[str] = None
+
+def generate_data_prompt(llm_commands: List[str], batch_size: int) -> str:
+    columns_description = "\n".join([
+        f"Column {i+1}: {cmd}" for i, cmd in enumerate(llm_commands)
+    ])
+    
+    return f"""Generate {batch_size} unique random rows of data based on these specifications:
+{columns_description}
+
+Requirements:
+- Each row must be different and realistic
+- Return ONLY a JSON array format: [["value1","value2"],["value1","value2"],...]
+- No additional text, explanations, or formatting
+- Values should be diverse and not repetitive
+
+JSON Array:"""
+
+@app.post("/generate", response_model=GenerationResponse)
+async def generate_data(request: GenerationRequest):
+    try:
+        # Set seed for reproducibility if provided
+        if request.seed:
+            torch.manual_seed(request.seed)
+            random.seed(request.seed)
+        
+        # Build prompt
+        prompt = generate_data_prompt(request.llm_commands, request.batch_size)
+        
+        # Prepare messages for chat template
+        messages = [
+            {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant that generates structured data."},
+            {"role": "user", "content": prompt}
+        ]
+        
+        # Apply chat template
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        
+        # Tokenize and generate
+        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            generated_ids = model.generate(
+                **model_inputs,
+                max_new_tokens=2048,
+                temperature=0.8,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode response
+        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]
+        
+        # Parse JSON from response
+        try:
+            # Find JSON array in the response
+            start_idx = response_text.find('[')
+            end_idx = response_text.rfind(']') + 1
+            
+            if start_idx == -1 or end_idx == 0:
+                raise ValueError("No JSON array found in response")
+            
+            json_str = response_text[start_idx:end_idx]
+            parsed_data = json.loads(json_str)
+            
+            # Validate data structure
+            if not isinstance(parsed_data, list):
+                raise ValueError("Response is not a list")
+            
+            # Filter and validate rows
+            valid_rows = []
+            expected_columns = len(request.llm_commands)
+            
+            for row in parsed_data:
+                if isinstance(row, list) and len(row) == expected_columns:
+                    # Convert all values to strings
+                    valid_rows.append([str(cell) for cell in row])
+            
+            return GenerationResponse(
+                success=True,
+                data=valid_rows
+            )
+            
+        except json.JSONDecodeError as e:
+            return GenerationResponse(
+                success=False,
+                data=[],
+                error=f"Failed to parse JSON: {str(e)}"
+            )
+        except Exception as e:
+            return GenerationResponse(
+                success=False,
+                data=[],
+                error=f"Data processing error: {str(e)}"
+            )
+            
+    except Exception as e:
+        return GenerationResponse(
+            success=False,
+            data=[],
+            error=f"Generation error: {str(e)}"
+        )
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "model": model_name}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)