Spaces:

griddava
/

pull-request-validator

Sleeping

App Files Files Community

Sgridda commited on Jul 11, 2025

Commit

733f0e1

1 Parent(s): a1f54c5

modified

Browse files

Files changed (2) hide show

main.py +62 -7
main_ai_version.py +161 -0

main.py CHANGED Viewed

@@ -4,6 +4,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import re
 import json
 # ----------------------------
 # 1. Configuration
@@ -76,12 +78,26 @@ def run_ai_inference(diff: str) -> str:
     if not model or not tokenizer:
         raise RuntimeError("Model is not loaded.")
-    # Prompt for codegen-350M-mono
-    prompt = f"""# Review this code and suggest improvements:\n{diff[:800]}\n# Review:"""
-    inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=1024, truncation=True)
     with torch.no_grad():
         outputs = model.generate(
-            inputs,
             max_new_tokens=128,
             do_sample=True,
             temperature=0.7,
@@ -91,10 +107,12 @@ def run_ai_inference(diff: str) -> str:
             eos_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id,
             use_cache=True
         )
-    response_text = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
-    # Post-process: take only the first non-empty line as the review
     review_lines = [line.strip() for line in response_text.strip().split('\n') if line.strip()]
-    review = review_lines[0] if review_lines else "AI review completed - no specific issues found."
     return review
 def parse_ai_response(response_text: str) -> list[ReviewComment]:
@@ -139,6 +157,43 @@ async def get_code_review(request: ReviewRequest):
 # ----------------------------
 # 7. Health Check Endpoint
 # ----------------------------
 @app.get("/health")
 async def health_check():

 import torch
 import re
 import json
+from fastapi.responses import HTMLResponse
 # ----------------------------
 # 1. Configuration
     if not model or not tokenizer:
         raise RuntimeError("Model is not loaded.")
+    # Improved prompt for codegen-350M-mono
+    prompt = (
+        "Below is a Python function. Please provide a code review comment with suggestions for improvement, in natural language. "
+        "Do not repeat the code.\n"
+        f"{diff[:800]}\n"
+        "Review comment:"
+    )
+    encoded = tokenizer(
+        prompt,
+        return_tensors="pt",
+        max_length=1024,
+        truncation=True,
+        padding="max_length"
+    )
+    input_ids = encoded["input_ids"]
+    attention_mask = encoded["attention_mask"]
     with torch.no_grad():
         outputs = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
             max_new_tokens=128,
             do_sample=True,
             temperature=0.7,
             eos_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id,
             use_cache=True
         )
+    response_text = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)
+    # Post-process: filter out code-like lines and fallback if needed
     review_lines = [line.strip() for line in response_text.strip().split('\n') if line.strip()]
+    # Filter out lines that look like code
+    comment_lines = [l for l in review_lines if not l.startswith("def ") and not l.startswith("class ") and not l.endswith(":") and not l.startswith("#")]
+    review = comment_lines[0] if comment_lines else "Consider adding a docstring and input validation."
     return review
 def parse_ai_response(response_text: str) -> list[ReviewComment]:
 # ----------------------------
 # 7. Health Check Endpoint
 # ----------------------------
+@app.get("/", response_class=HTMLResponse)
+def root_html():
+    """Return HTML for browser viewing."""
+    return """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>AI Code Review Service</title>
+        <style>
+            body { font-family: Arial, sans-serif; margin: 40px; }
+            .status { color: green; font-weight: bold; }
+            .endpoint { background: #f4f4f4; padding: 10px; margin: 10px 0; border-radius: 5px; }
+        </style>
+    </head>
+    <body>
+        <h1>AI Code Review Service</h1>
+        <p class="status">✅ Service is running in emergency mode</p>
+        <h2>Available Endpoints:</h2>
+        <div class="endpoint"><strong>GET /health</strong> - Health check</div>
+        <div class="endpoint"><strong>POST /review</strong> - Submit code diff for review</div>
+        <div class="endpoint"><strong>GET /docs</strong> - Interactive API documentation</div>
+        <div class="endpoint"><strong>GET /test</strong> - Simple test endpoint</div>
+        <h2>Quick Test:</h2>
+        <p><a href="/health">Test Health Endpoint</a></p>
+        <p><a href="/docs">View API Documentation</a></p>
+        <h2>Status:</h2>
+        <ul>
+            <li>Mode: Emergency (Mock responses)</li>
+            <li>AI Model: Disabled</li>
+            <li>Response Time: ~100ms</li>
+        </ul>
+    </body>
+    </html>
+    """
 @app.get("/health")
 async def health_check():

main_ai_version.py ADDED Viewed

	@@ -0,0 +1,161 @@

+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import re
+import json
+# ----------------------------
+# 1. Configuration
+# ----------------------------
+MODEL_NAME = "Salesforce/codegen-350M-mono"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# ----------------------------
+# 2. FastAPI App Initialization
+# ----------------------------
+app = FastAPI(
+    title="AI Code Review Service",
+    description="An API to get AI-powered code reviews for pull request diffs.",
+    version="1.0.0",
+)
+# ----------------------------
+# 3. AI Model Loading
+# ----------------------------
+model = None
+tokenizer = None
+def load_model():
+    """Loads the model and tokenizer into memory."""
+    global model, tokenizer
+    if model is None:
+        print(f"Loading model: {MODEL_NAME} on device: {DEVICE}...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.float32,
+            device_map="cpu",
+        )
+        print("Model loaded successfully.")
+@app.on_event("startup")
+async def startup_event():
+    """
+    On server startup, we trigger the model loading.
+    """
+    print("Server starting up...")
+    load_model()
+# ----------------------------
+# 4. API Request/Response Models
+# ----------------------------
+class ReviewRequest(BaseModel):
+    diff: str
+class ReviewComment(BaseModel):
+    file_path: str
+    line_number: int
+    comment_text: str
+class ReviewResponse(BaseModel):
+    comments: list[ReviewComment]
+# ----------------------------
+# 5. The AI Review Logic
+# ----------------------------
+def run_ai_inference(diff: str) -> str:
+    """
+    Runs the AI model to get the review.
+    """
+    if not model or not tokenizer:
+        raise RuntimeError("Model is not loaded.")
+    # Improved prompt for codegen-350M-mono
+    prompt = (
+        "Below is a Python function. Please provide a code review comment with suggestions for improvement, in natural language. "
+        "Do not repeat the code.\n"
+        f"{diff[:800]}\n"
+        "Review comment:"
+    )
+    encoded = tokenizer(
+        prompt,
+        return_tensors="pt",
+        max_length=1024,
+        truncation=True,
+        padding="max_length"
+    )
+    input_ids = encoded["input_ids"]
+    attention_mask = encoded["attention_mask"]
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=128,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.95,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id,
+            use_cache=True
+        )
+    response_text = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)
+    # Post-process: filter out code-like lines and fallback if needed
+    review_lines = [line.strip() for line in response_text.strip().split('\n') if line.strip()]
+    # Filter out lines that look like code
+    comment_lines = [l for l in review_lines if not l.startswith("def ") and not l.startswith("class ") and not l.endswith(":") and not l.startswith("#")]
+    review = comment_lines[0] if comment_lines else "Consider adding a docstring and input validation."
+    return review
+def parse_ai_response(response_text: str) -> list[ReviewComment]:
+    """
+    Parses the raw text from the AI to extract the JSON array.
+    """
+    # For codegen-350M-mono, just wrap the review in a single comment
+    return [ReviewComment(
+        file_path="code_reviewed.py",
+        line_number=1,
+        comment_text=response_text.strip()
+    )]
+# ----------------------------
+# 6. The API Endpoint
+# ----------------------------
+@app.post("/review", response_model=ReviewResponse)
+async def get_code_review(request: ReviewRequest):
+    if not request.diff:
+        raise HTTPException(status_code=400, detail="Diff content cannot be empty.")
+    import time
+    start_time = time.time()
+    print(f"Starting review request at {start_time}")
+    try:
+        print("Running AI inference...")
+        ai_response_text = run_ai_inference(request.diff)
+        print(f"AI inference completed in {time.time() - start_time:.2f} seconds")
+        print("Parsing AI response...")
+        parsed_comments = parse_ai_response(ai_response_text)
+        print(f"Total processing time: {time.time() - start_time:.2f} seconds")
+        return ReviewResponse(comments=parsed_comments)
+    except Exception as e:
+        print(f"An unexpected error occurred after {time.time() - start_time:.2f} seconds: {e}")
+        raise HTTPException(status_code=500, detail="An internal error occurred while processing the review.")
+# ----------------------------
+# 7. Health Check Endpoint
+# ----------------------------
+@app.get("/health")
+async def health_check():
+    return {"status": "ok", "model_loaded": model is not None}