feat: Introduce FastAPI endpoints for single and batch text generation with Pydantic models and Hugging Face model management.

models/generate.py

@@ -1,14 +1,52 @@
 """Pydantic models for text generation"""
 
-from pydantic import BaseModel
-from typing import Optional
+from pydantic import BaseModel, ConfigDict, Field
+from typing import Optional, List
 
 class GenerateRequest(BaseModel):
-    input_text: str
-    max_length: Optional[int] = 128
-    num_beams: Optional[int] = 4
+    """Request model for argument generation"""
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "topic": "Assisted suicide should be a criminal offence",
+                "position": "positive"  # "positive" or "negative"
+            }
+        }
+    )
+    
+    topic: str = Field(..., min_length=5, max_length=1000, 
+                       description="The debate topic or statement")
+    position: str = Field(..., min_length=5, max_length=50,
+                         description="The stance to take")
+
 
 class GenerateResponse(BaseModel):
-    input_text: str
-    generated_text: str
+    """Response model for argument generation"""
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "topic": "Assisted suicide should be a criminal offence",
+                "position": "positive",  # "positive" or "negative"
+                "argument": "People have the right to choose how they end their lives",
+                "timestamp": "2024-11-15T10:30:00"
+            }
+        }
+    )
+    
+    topic: str  
+    position: str
+    argument: str
+    timestamp: str
+    timestamp: str
+
+
+class BatchGenerateRequest(BaseModel):
+    """Request model for batch argument generation"""
+    items: List[GenerateRequest]
+
+
+class BatchGenerateResponse(BaseModel):
+    """Response model for batch argument generation"""
+    results: List[GenerateResponse]
+    model_info: Optional[str] = "KPA T5 Generation Model"
     timestamp: str

routes/generate.py

@@ -5,41 +5,75 @@ from datetime import datetime
 import logging
 
 from services import generate_model_manager
-from models.generate import GenerateRequest, GenerateResponse
+from models.generate import GenerateRequest, GenerateResponse, BatchGenerateRequest, BatchGenerateResponse
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
 
 
-@router.post("/generate", response_model=GenerateResponse, tags=["Text Generation"])
-async def generate_text(request: GenerateRequest):
+@router.post("/predict", response_model=GenerateResponse, tags=["Text Generation"])
+async def generate_argument(request: GenerateRequest):
     """
-    Generate text using the T5 model
+    Generate an argument for a given topic and position
     
-    - **input_text**: The input text for generation
-    - **max_length**: Maximum length of generated text (default: 128)
-    - **num_beams**: Number of beams for beam search (default: 4)
-    
-    Returns generated text
+    - **topic**: The debate topic
+    - **position**: The stance (e.g. "positive", "negative")
     """
     try:
         # Generate text
         result = generate_model_manager.generate(
-            request.input_text, 
-            max_length=request.max_length,
-            num_beams=request.num_beams
+            topic=request.topic,
+            position=request.position
         )
         
         # Build response
         response = GenerateResponse(
-            input_text=request.input_text,
-            generated_text=result,
+            topic=request.topic,
+            position=request.position,
+            argument=result,
             timestamp=datetime.now().isoformat()
         )
         
-        logger.info(f"Generated text: {result}")
+        logger.info(f"Generated argument: {result[:50]}...")
         return response
         
     except Exception as e:
         logger.error(f"Generation error: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
+
+
+@router.post("/batch-predict", response_model=BatchGenerateResponse, tags=["Text Generation"])
+async def batch_generate_argument(request: BatchGenerateRequest):
+    """
+    Generate arguments for multiple topic-position pairs
+    """
+    try:
+        items_data = [{"topic": item.topic, "position": item.position} for item in request.items]
+        
+        # Batch generate
+        results = generate_model_manager.batch_generate(
+            items=items_data
+        )
+        
+        # Build response
+        response_items = []
+        timestamp = datetime.now().isoformat()
+        
+        for i, item in enumerate(request.items):
+            response_items.append(
+                GenerateResponse(
+                    topic=item.topic,
+                    position=item.position,
+                    argument=results[i],
+                    timestamp=timestamp
+                )
+            )
+            
+        return BatchGenerateResponse(
+            results=response_items,
+            timestamp=timestamp
+        )
+
+    except Exception as e:
+        logger.error(f"Batch generation error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Batch generation failed: {str(e)}")

services/generate_model_manager.py

@@ -56,11 +56,18 @@ class GenerateModelManager:
             logger.error(f"Error loading generation model: {str(e)}")
             raise RuntimeError(f"Failed to load generation model: {str(e)}")
     
-    def generate(self, input_text: str, max_length: int = 128, num_beams: int = 4) -> str:
-        """Generate text from input"""
+    def _format_input(self, topic: str, position: str) -> str:
+        """Format input for the model"""
+        # Standard format for argument generation
+        return f"topic: {topic} stance: {position}"
+
+    def generate(self, topic: str, position: str, max_length: int = 128, num_beams: int = 4) -> str:
+        """Generate argument for a topic and position"""
         if not self.model_loaded:
             raise RuntimeError("Generation model not loaded")
         
+        input_text = self._format_input(topic, position)
+        
         # Tokenize
         inputs = self.tokenizer(
             input_text,
@@ -84,6 +91,37 @@ class GenerateModelManager:
         
         return generated_text
 
+    def batch_generate(self, items: list[dict], max_length: int = 128, num_beams: int = 4) -> list[str]:
+        """Batch generate arguments"""
+        if not self.model_loaded:
+            raise RuntimeError("Generation model not loaded")
+        
+        # Prepare inputs
+        input_texts = [self._format_input(item["topic"], item["position"]) for item in items]
+        
+        # Tokenize batch
+        inputs = self.tokenizer(
+            input_texts,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512,
+            padding=True
+        ).to(self.device)
+        
+        # Generate batch
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_length=max_length,
+                num_beams=num_beams,
+                early_stopping=True
+            )
+            
+        # Decode batch
+        generated_texts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        
+        return generated_texts
+
 
 # Initialize singleton instance
 generate_model_manager = GenerateModelManager()