added inferernce

main.py

@@ -5,6 +5,7 @@ import torch
 import re
 import json
 from fastapi.responses import HTMLResponse
+import requests
 
 
 # ----------------------------
@@ -14,6 +15,10 @@ from fastapi.responses import HTMLResponse
 MODEL_NAME = "Salesforce/codegen-350M-mono"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Add your Hugging Face API key here
+HF_API_KEY = "REDACTED"
+HF_MODEL_NAME = "bigcode/starcoder"  # Replace with the best model for code review
+
 # ----------------------------
 # 2. FastAPI App Initialization
 # ----------------------------
@@ -75,61 +80,76 @@ class ReviewResponse(BaseModel):
 
 def run_ai_inference(diff: str) -> str:
     """
-    Runs the AI model to get the review.
+    Sends the code diff to Hugging Face Inference API to get the review.
     """
-    if not model or not tokenizer:
-        raise RuntimeError("Model is not loaded.")
-
-    # Simple, direct prompt for codegen-350M-mono
-    prompt = f"""Code:
-{diff[:500]}
-
-Review: This code could be improved by adding"""
-    encoded = tokenizer(
-        prompt,
-        return_tensors="pt",
-        max_length=512,  # Reduced from 1024 for faster processing
-        truncation=True,
-        padding="max_length"
+    if not HF_API_KEY:
+        raise RuntimeError("Hugging Face API key is not set.")
+
+    # Better prompt for meaningful completions
+    prompt = f"""# Code Review
+
+def example():
+    pass
+# Review: This function should include error handling and documentation.
+
+{diff[:400]}
+# Review: This code should include"""
+
+    headers = {
+        "Authorization": f"Bearer {HF_API_KEY}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": 32,
+            "temperature": 0.7,
+            "top_p": 0.9
+        }
+    }
+
+    response = requests.post(
+        f"https://api-inference.huggingface.co/models/{HF_MODEL_NAME}",
+        headers=headers,
+        json=payload
     )
-    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=32,  # Further reduced for speed
-            do_sample=True,
-            temperature=0.9,
-            top_p=0.85,
-            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)
-    
+
+    if response.status_code != 200:
+        raise RuntimeError(f"Hugging Face API error: {response.status_code} {response.text}")
+
+    response_data = response.json()
+    if isinstance(response_data, list) and len(response_data) > 0:
+        response_text = response_data[0].get("generated_text", "").strip()
+    else:
+        response_text = "Unable to generate a meaningful review."
+
     # Clean up the response
     response_text = response_text.strip()
-    
-    # Remove artifacts and clean up
-    if response_text.startswith("adding"):
-        response_text = "Adding " + response_text[6:]
-    
-    # Take only the first sentence or meaningful phrase
-    sentences = response_text.split('.')
-    if sentences and len(sentences[0].strip()) > 10:
-        review = sentences[0].strip() + "."
+
+    # Handle different completion patterns
+    if response_text.startswith("error handling"):
+        review = "Consider adding error handling and input validation."
+    elif response_text.startswith("documentation"):
+        review = "Consider adding documentation and type hints."
+    elif response_text.startswith("input validation"):
+        review = "Consider adding input validation and error checks."
+    elif response_text.startswith("type hints"):
+        review = "Consider adding type hints and documentation."
     else:
-        # Fallback to first meaningful line
+        # Extract meaningful content
         lines = [line.strip() for line in response_text.split('\n') if line.strip()]
-        if lines and len(lines[0]) > 5:
-            review = lines[0]
-            if not review.endswith('.'):
-                review += "."
+        if lines and len(lines[0]) > 3:
+            first_line = lines[0]
+            # Clean up common artifacts
+            if first_line.startswith('#'):
+                first_line = first_line[1:].strip()
+            if len(first_line) > 10:
+                review = f"Consider adding {first_line.lower()}."
+            else:
+                review = "Consider adding proper documentation and error handling."
         else:
             review = "Consider adding proper documentation and error handling."
-    
+
     return review
 
 def parse_ai_response(response_text: str) -> list[ReviewComment]: