from fastapi import FastAPI
from transformers import pipeline
import json

app = FastAPI()

# Carregando o MedGemma direto do Hugging Face
medgemma = pipeline(
    "text-generation",
    model="google/medgemma-4b-it",
    tokenizer="google/medgemma-4b-it"
)

@app.get("/generate")
async def generate_question(theme: str, difficulty: str):
    prompt = f"Gere uma questão sobre {theme}, nível {difficulty}, com 4 alternativas (A, B, C, D), no estilo da USP."
    response = medgemma(prompt, max_new_tokens=300, temperature=0.7, top_p=0.95)[0]['generated_text']
    return {"question": response}

@app.get("/simulado")
async def get_simulado(num_questions: int = 5):
    simulado = []
    for _ in range(num_questions):
        prompt = "Gere uma questão médica estilo USP com 4 alternativas (A, B, C, D), nível médio."
        question = medgemma(prompt, max_new_tokens=300)[0]['generated_text']
        simulado.append({"question": question})
    return {"simulado": simulado}