made it simple

main.py

@@ -99,12 +99,17 @@ def run_ai_inference(diff: str) -> str:
 
     inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
     
+    # Create attention mask to avoid warnings and improve reliability
+    attention_mask = torch.ones_like(inputs)
+    
     # Optimized generation parameters for speed
     outputs = model.generate(
-        inputs, 
-        max_new_tokens=256,  # Reduced from 1024
-        do_sample=False,
-        temperature=0.1,     # Lower temperature for more focused output
+        inputs,
+        attention_mask=attention_mask,
+        max_new_tokens=128,  # Further reduced for faster generation
+        do_sample=True,      # Enable sampling to use temperature
+        temperature=0.3,     # Lower temperature for more focused output
+        top_p=0.9,           # Nucleus sampling for better quality
         num_return_sequences=1, 
         eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.eos_token_id,

main_lightweight.py

@@ -0,0 +1,139 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import torch
+import logging
+import json
+import re
+
+# Ultra-lightweight version with minimal AI
+app = FastAPI(
+    title="AI Code Review Service",
+    description="An API to get AI-powered code reviews for pull request diffs.",
+    version="1.0.0",
+)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Try to load a very small model, fall back to mock if it fails
+model = None
+tokenizer = None
+
+def load_simple_model():
+    """Try to load the smallest possible model."""
+    global model, tokenizer
+    try:
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        
+        # Use the smallest possible model
+        model_name = "distilgpt2"  # Much smaller than TinyLlama
+        logger.info("Loading lightweight model: %s", model_name)
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+        )
+        logger.info("Model loaded successfully")
+        return True
+    except Exception as e:
+        logger.warning("Failed to load AI model: %s. Using mock responses.", str(e))
+        return False
+
+# Try to load model on startup
+model_loaded = load_simple_model()
+
+class DiffRequest(BaseModel):
+    diff: str
+
+class ReviewComment(BaseModel):
+    file_path: str
+    line_number: int
+    comment_text: str
+
+class ReviewResponse(BaseModel):
+    comments: list[ReviewComment]
+
+@app.get("/health")
+def health_check():
+    """Health check endpoint."""
+    return {
+        "status": "healthy",
+        "service": "AI Code Review Service",
+        "model_loaded": model_loaded,
+        "model_name": "distilgpt2" if model_loaded else "mock",
+        "device": "cuda" if torch.cuda.is_available() else "cpu"
+    }
+
+def simple_ai_review(diff: str):
+    """Very simple AI review using the lightweight model."""
+    if not model_loaded or not model or not tokenizer:
+        return None
+        
+    try:
+        # Very simple prompt
+        prompt = f"Review this code change and suggest improvements:\n{diff[:200]}\nSuggestion:"
+        
+        inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=256, truncation=True)
+        
+        # Very conservative generation
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_new_tokens=50,  # Very short response
+                do_sample=False,
+                num_return_sequences=1,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=True
+            )
+        
+        response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+        return response.strip()
+    except Exception as e:
+        logger.warning("AI generation failed: %s", str(e))
+        return None
+
+@app.post("/review", response_model=ReviewResponse)
+def review_diff(request: DiffRequest):
+    """Review endpoint with fallback to mock data."""
+    logger.info("Received diff for review (length: %d chars)", len(request.diff))
+    
+    # Try AI first, fall back to mock
+    ai_suggestion = None
+    if model_loaded:
+        ai_suggestion = simple_ai_review(request.diff)
+    
+    if ai_suggestion:
+        # Use AI suggestion
+        comments = [{
+            "file_path": "reviewed_file.py",
+            "line_number": 1,
+            "comment_text": ai_suggestion
+        }]
+        logger.info("Returning AI-generated review")
+    else:
+        # Fall back to mock comments
+        comments = [
+            {
+                "file_path": "example.py",
+                "line_number": 1,
+                "comment_text": "Consider adding error handling and input validation."
+            },
+            {
+                "file_path": "example.py",
+                "line_number": 5,
+                "comment_text": "This function could benefit from better documentation."
+            }
+        ]
+        logger.info("Returning mock review comments")
+    
+    return ReviewResponse(comments=comments)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

main_simple.py

@@ -0,0 +1,66 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import json
+import logging
+
+# Simple version without AI model for testing
+app = FastAPI(
+    title="AI Code Review Service",
+    description="An API to get AI-powered code reviews for pull request diffs.",
+    version="1.0.0",
+)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class DiffRequest(BaseModel):
+    diff: str
+
+class ReviewComment(BaseModel):
+    file_path: str
+    line_number: int
+    comment_text: str
+
+class ReviewResponse(BaseModel):
+    comments: list[ReviewComment]
+
+@app.get("/health")
+def health_check():
+    """Health check endpoint."""
+    return {
+        "status": "healthy",
+        "service": "AI Code Review Service",
+        "model_loaded": False,  # No model in simple version
+        "message": "Simple version - returns mock reviews"
+    }
+
+@app.post("/review", response_model=ReviewResponse)
+def review_diff(request: DiffRequest):
+    """
+    Mock review endpoint that returns sample comments.
+    Replace this with actual AI logic once the Space is working.
+    """
+    logger.info("Received diff for review (length: %d chars)", len(request.diff))
+    
+    # Mock review comments
+    mock_comments = [
+        {
+            "file_path": "example.py",
+            "line_number": 1,
+            "comment_text": "Consider adding docstrings to improve code documentation."
+        },
+        {
+            "file_path": "example.py", 
+            "line_number": 5,
+            "comment_text": "This function could benefit from error handling."
+        }
+    ]
+    
+    logger.info("Returning %d mock review comments", len(mock_comments))
+    
+    return ReviewResponse(comments=mock_comments)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)