shekkari21/codereviewer

1from fastapi import FastAPI, Request, Form2from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse3from pydantic import BaseModel4from typing import List5from clearml import Model6import torch7from configs import add_args8from models import build_or_load_gen_model9import argparse10from argparse import Namespace11import os12from peft import PeftModel, PeftConfig, get_peft_model, LoraConfig13 14MAX_SOURCE_LENGTH = 51215 16def pad_assert(tokenizer, source_ids):17    source_ids = source_ids[:MAX_SOURCE_LENGTH - 2]18    source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]19    pad_len = MAX_SOURCE_LENGTH - len(source_ids)20    source_ids += [tokenizer.pad_id] * pad_len21    assert len(source_ids) == MAX_SOURCE_LENGTH, "Not equal length."22    return source_ids23 24# Encode code content and comment into model input25def encode_diff(tokenizer, code, comment):26    # Tokenize code file content27    code_ids = tokenizer.encode(code, max_length=MAX_SOURCE_LENGTH, truncation=True)[1:-1]28    # Tokenize comment29    comment_ids = tokenizer.encode(comment, max_length=MAX_SOURCE_LENGTH, truncation=True)[1:-1]30    # Concatenate: [BOS] + code + [EOS] + [msg_id] + comment31    source_ids = [tokenizer.bos_id] + code_ids + [tokenizer.eos_id]32    source_ids += [tokenizer.msg_id] + comment_ids33    # Pad/truncate to fixed length34    source_ids = source_ids[:MAX_SOURCE_LENGTH - 2]35    source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]36    pad_len = MAX_SOURCE_LENGTH - len(source_ids)37    source_ids += [tokenizer.pad_id] * pad_len38    assert len(source_ids) == MAX_SOURCE_LENGTH, "Not equal length."39    return source_ids40 41# Load base model architecture and tokenizer from HuggingFace42BASE_MODEL_NAME = "microsoft/codereviewer"43args = Namespace(44    model_name_or_path=BASE_MODEL_NAME,45    load_model_path=None,46    # Add other necessary default arguments if build_or_load_gen_model requires them47)48print(f"Loading base model architecture and tokenizer from: {BASE_MODEL_NAME}")49config, base_model, tokenizer = build_or_load_gen_model(args)50print("Base model architecture and tokenizer loaded.")51 52# Download the fine-tuned weights from ClearML53CLEARML_MODEL_ID = "34e25deb24c64b74b29c8519ed15fe3e"54model_obj = Model(model_id=CLEARML_MODEL_ID)55finetuned_weights_path = model_obj.get_local_copy()56adapter_dir = os.path.dirname(finetuned_weights_path)57 58print(f"Fine-tuned adapter weights downloaded to directory: {adapter_dir}")59 60# Create LoRA configuration matching the fine-tuned checkpoint61lora_cfg = LoraConfig(62    r=64,63    lora_alpha=128,64    target_modules=["q", "wo", "wi", "v", "o", "k"],65    lora_dropout=0.05,66    bias="none",67    task_type="SEQ_2_SEQ_LM"68)69# Wrap base model with PEFT LoRA70peft_model = get_peft_model(base_model, lora_cfg)71# Load adapter-only weights and merge into base72adapter_state = torch.load(finetuned_weights_path, map_location="cpu")73peft_model.load_state_dict(adapter_state, strict=False)74model = peft_model.merge_and_unload()75print("Merged base model with LoRA adapters.")76 77model.to("cpu")78model.eval()79print("Model ready for inference.")80 81app = FastAPI()82 83last_payload = {"comment": "", "files": []}84last_infer_result = {"generated_code": ""}85 86class FileContent(BaseModel):87    filename: str88    content: str89 90class PRPayload(BaseModel):91    comment: str92    files: List[FileContent]93 94class InferenceRequest(BaseModel):95    comment: str96    files: List[FileContent]97 98 99@app.get("/")100def root():101    return {"message": "FastAPI PR comment service is running"}102 103@app.post("/pr-comments")104async def receive_pr_comment(payload: PRPayload):105    global last_payload106    last_payload = payload.dict()107    # Return the received payload as JSON and also redirect to /show108    return JSONResponse(content={"status": "received", "payload": last_payload, "redirect": "/show"})109 110@app.get("/show", response_class=HTMLResponse)111def show_last_comment():112    html = f"<h2>Received Comment</h2><p>{last_payload['comment']}</p><hr>"113    for file in last_payload["files"]:114        html += f"<h3>{file['filename']}</h3><pre>{file['content']}</pre><hr>"115    return html116 117@app.post("/infer")118async def infer(request: InferenceRequest):119    global last_infer_result120    print("[DEBUG] Received /infer request with:", request.dict())121    122    code = request.files[0].content if request.files else ""123    source_ids = encode_diff(tokenizer, code, request.comment)124    # print("[DEBUG] source_ids:", source_ids)125    #tokens = [tokenizer.decode([sid], skip_special_tokens=False) for sid in source_ids]126    #print("[DEBUG] tokens:", tokens)127    inputs = torch.tensor([source_ids], dtype=torch.long)128    inputs_mask = inputs.ne(tokenizer.pad_id)129    130    preds = model.generate(131        inputs,132        attention_mask=inputs_mask,133        use_cache=True,134        num_beams=5,135        early_stopping=True,136        max_length=100,137        num_return_sequences=1138    )139    140    pred = preds[0].cpu().numpy()141    pred_nl = tokenizer.decode(pred[2:], skip_special_tokens=True, clean_up_tokenization_spaces=False)142    last_infer_result = {"generated_code": pred_nl}143    return last_infer_result144 145@app.get("/show-infer", response_class=HTMLResponse)146def show_infer_result():147    html = f"<h2>Generated Message</h2><pre>{last_infer_result['generated_code']}</pre>"148    return html149 150if __name__ == "__main__":151    # Place any CLI/training logic here if needed152    # This block is NOT executed when running with uvicorn153    pass