tinyllama_inference.py

# CODEXGAME/backend/ai_evaluator/tinyllama_inference.py

import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def load_model():
    # Change the model identifier if needed – this should be a TinyLlama variant available on Hugging Face.
    model_name = "TheBloke/tiny-llama-7b"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return tokenizer, model

def evaluate_code(question, code):
    # Construct a prompt for the AI evaluator.
    prompt = f"""

You are an expert code evaluator.

Rate the user's solution to the following problem from 0-5 (0 = completely incorrect, 5 = excellent).

Also provide a concise "feedback" message.

Problem: "{question}"

Solution: "{code}"

Return ONLY valid JSON: {{"stars": number, "feedback": string}}

Do not include any extra text outside the JSON.

"""
    tokenizer, model = load_model()
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=150)
    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    try:
        result = json.loads(response_text.strip())
    except Exception as e:
        result = {"stars": 0, "feedback": "Evaluation failed. Unable to parse AI response."}
    return result

# For direct testing from the command line
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))