File size: 2,008 Bytes
e65d0ad
0bae633
e65d0ad
 
3d71f2e
 
 
e65d0ad
3d71f2e
 
 
 
 
e65d0ad
 
 
3d71f2e
0bae633
3d71f2e
 
 
 
 
e65d0ad
 
 
 
 
0bae633
 
3d71f2e
 
 
 
0bae633
e65d0ad
0bae633
 
 
 
 
 
 
 
 
 
 
e65d0ad
 
 
 
 
 
0bae633
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import json
import re
from transformers import AutoTokenizer, AutoModelForCausalLM

# Global variables to cache the model and tokenizer
tokenizer, model = None, None

def load_model():
    global tokenizer, model
    if tokenizer is None or model is None:
        model_name = "Salesforce/codegen-350M-mono"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)
    return tokenizer, model

def evaluate_code(question, code):
    # Refined prompt instructing the model to output exactly valid JSON.
    prompt = f"""You are an expert code evaluator.
Evaluate the following solution for the given problem.
Return ONLY a JSON object with exactly two keys:
  "stars": an integer between 0 and 5 (0 means completely incorrect, 5 means excellent).
  "feedback": a concise message.
Output must be exactly valid JSON and nothing else.
Problem: "{question}"
Solution: "{code}"
"""
    tokenizer, model = load_model()
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        temperature=0.0,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=False
    )
    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    match = re.search(r'\{.*\}', response_text)
    if match:
        json_text = match.group(0)
        try:
            result = json.loads(json_text)
        except Exception as e:
            result = {"stars": 0, "feedback": "Evaluation failed. Unable to parse AI response."}
    else:
        result = {"stars": 0, "feedback": "Evaluation failed. Unable to extract JSON from AI response."}
    
    return result

if __name__ == "__main__":
    import sys
    if len(sys.argv) < 3:
        print(json.dumps({"error": "Please provide a question and code as arguments"}))
        sys.exit(1)
    question = sys.argv[1]
    code = sys.argv[2]
    result = evaluate_code(question, code)
    print(json.dumps(result))