from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel import torch from transformers import AutoTokenizer, AutoModelForCausalLM import random import re from typing import List, Dict, Any, Optional app = FastAPI(title="CodeGen Kids Tutor API") # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], # For production, specify your frontend domain allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Model loading print("Loading model and tokenizer...") MODEL_NAME = "AhmedMOstaFA10/codegen-kids-tutor" tokenizer = None model = None class ProblemRequest(BaseModel): category: Optional[str] = None # Optional category to filter problem prompts class SolutionRequest(BaseModel): code: str reference_code: str # Problem prompts categorized by topic problem_prompts = { "arithmetic": [ "# Instruction:\nGenerate a simple arithmetic problem suitable for a kid. Write a function with a short docstring and partial code.\n\n" "# Input:\nAddition, subtraction, or multiplication\n\n" "# Solution:\n" ], "strings": [ "# Instruction:\nGenerate a basic string manipulation exercise suitable for a beginner. Write a function with a short docstring and partial code.\n\n" "# Input:\nA string operation like reversing, counting characters, or checking substrings\n\n" "# Solution:\n" ], "lists": [ "# Instruction:\nGenerate a simple list-related problem for beginners. Write a function with a short docstring and partial implementation.\n\n" "# Input:\nSorting a list, finding max or min, or summing numbers\n\n" "# Solution:\n" ], "conditions": [ "# Instruction:\nGenerate a basic Python problem using if-else conditions. Write a function with a docstring and a few lines of partial code.\n\n" "# Input:\nAge check, number comparison, or grade classification\n\n" "# Solution:\n" ], "loops": [ "# Instruction:\nCreate a beginner-friendly problem that uses a for loop. Write a function with a clear docstring and partial implementation.\n\n" "# Input:\nSumming numbers, iterating over lists, or counting even numbers\n\n" "# Solution:\n", "# Instruction:\nWrite a basic programming problem involving a while loop. Include a function definition, a short docstring, and partial implementation.\n\n" "# Input:\nRepeating until condition is met, counting, or basic input validation\n\n" "# Solution:\n" ], "dictionaries": [ "# Instruction:\nGenerate an easy dictionary-based Python exercise. Write a function with a short docstring and partial implementation.\n\n" "# Input:\nAccessing values, summing values, or checking keys in a dictionary\n\n" "# Solution:\n" ], "input_output": [ "# Instruction:\nWrite a problem simulating user input and output in Python. Provide a function with a docstring and a few lines of implementation.\n\n" "# Input:\nName, age, or favorite color, and return a formatted string\n\n" "# Solution:\n" ], "math": [ "# Instruction:\nGenerate a Python problem that implements a basic math formula. Include a function with a docstring and partial code.\n\n" "# Input:\nArea of circle, BMI calculation, or temperature conversion\n\n" "# Solution:\n" ], "boolean": [ "# Instruction:\nCreate a beginner-friendly Python exercise using boolean logic. Write a function with a docstring and partial implementation.\n\n" "# Input:\nCheck conditions like even AND positive, or NOT equal to zero\n\n" "# Solution:\n" ] } # Get all prompts in a single list for random selection all_prompts = [] for category_prompts in problem_prompts.values(): all_prompts.extend(category_prompts) @app.on_event("startup") async def startup_event(): global tokenizer, model try: tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) model.config.pad_token_id = tokenizer.pad_token_id # Check for GPU availability device = "cuda" if torch.cuda.is_available() else "cpu" model = model.to(device) print(f"Model loaded successfully on {device}") except Exception as e: print(f"Error loading model: {str(e)}") # We'll initialize lazily if this fails on startup def get_model(): global tokenizer, model if tokenizer is None or model is None: tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) model.config.pad_token_id = tokenizer.pad_token_id # Check for GPU availability device = "cuda" if torch.cuda.is_available() else "cpu" model = model.to(device) return tokenizer, model def generate_full_solution(prompt): tokenizer, model = get_model() inputs = tokenizer(prompt, return_tensors="pt").to(model.device) with torch.no_grad(): generated_ids = model.generate( inputs["input_ids"], max_length=256, num_return_sequences=1, temperature=0.7, top_p=0.95, do_sample=True, pad_token_id=tokenizer.pad_token_id ) full_solution = tokenizer.decode(generated_ids[0], skip_special_tokens=True) return full_solution def truncate_function_body(code): lines = code.strip().split('\n') truncated = [] for line in lines: stripped = line.strip() truncated.append(line) if stripped.startswith('return') or stripped.startswith('print'): break if len(truncated) >= 4: break return '\n'.join(truncated) @app.get("/") def read_root(): return {"message": "CodeGen Kids Tutor API is running!"} @app.post("/generate-problem") def generate_problem(request: ProblemRequest): try: # Select prompts based on category if provided selected_prompts = [] if request.category and request.category in problem_prompts: selected_prompts = problem_prompts[request.category] else: selected_prompts = all_prompts if not selected_prompts: raise HTTPException(status_code=400, detail="No problem prompts available for the selected category") problem_prompt = random.choice(selected_prompts) complete_solution = generate_full_solution(problem_prompt) # Extract problem statement and function code split = complete_solution.strip().split('\n') problem_lines = [] function_lines = [] for line in split: if line.strip().startswith("def ") or line.strip().startswith('"""') or line.strip().startswith("#"): function_lines.append(line) else: problem_lines.append(line) current_problem = '\n'.join(problem_lines[:2]).strip() truncated_solution = truncate_function_body('\n'.join(function_lines)) return { "problem": current_problem, "starter_code": truncated_solution, "reference_code": truncated_solution # For verification later } except Exception as e: raise HTTPException(status_code=500, detail=f"Error generating problem: {str(e)}") @app.post("/check-solution") def check_solution(request: SolutionRequest): try: user_solution = request.code.strip() reference_code = request.reference_code.strip() # Basic syntax check try: compile(user_solution, '', 'exec') except Exception as e: return { "is_correct": False, "feedback": f"Syntax error: {str(e)}" } # Function name check model_func_match = re.search(r'def\s+([a-zA-Z_][a-zA-Z0-9_]*)', reference_code) user_func_match = re.search(r'def\s+([a-zA-Z_][a-zA-Z0-9_]*)', user_solution) if model_func_match and user_func_match: if model_func_match.group(1) != user_func_match.group(1): return { "is_correct": False, "feedback": "You changed the function name. Keep the original function name." } # Import difflib for sequence matching to evaluate solution similarity from difflib import SequenceMatcher similarity = SequenceMatcher(None, reference_code, user_solution).ratio() if similarity > 0.5: return { "is_correct": True, "feedback": "Your solution looks correct! Great job! 🎉" } elif similarity > 0.3: return { "is_correct": True, "feedback": "Your solution passes, but there might be a more efficient approach. Keep going! 👍" } else: return { "is_correct": False, "feedback": "Your solution differs significantly from the expected solution. Try again! 🔄" } except Exception as e: raise HTTPException(status_code=500, detail=f"Error checking solution: {str(e)}")