import os
import torch
from fastapi import FastAPI
from transformers import AutoProcessor, AutoModelForCausalLM
from pydantic import BaseModel
import logging

# ログ設定
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# モデルロード
model_name = "google/gemma-3-4b-it"  # 軽量な2Bモデルに変更
try:
    logger.info(f"Loading model: {model_name}")
    processor = AutoProcessor.from_pretrained(model_name, token=os.getenv("HF_TOKEN"))
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        token=os.getenv("HF_TOKEN"),
        low_cpu_mem_usage=True,
        load_in_4bit=True  # 量子化でメモリ節約
    )
    logger.info("Model loaded successfully")
except Exception as e:
    logger.error(f"Model load error: {e}")
    raise

class TextInput(BaseModel):
    text: str
    max_length: int = 50

@app.post("/generate")
async def generate_text(input: TextInput):
    try:
        logger.info(f"Generating text for input: {input.text}")
        inputs = processor(input.text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
        outputs = model.generate(**inputs, max_length=input.max_length)
        result = processor.decode(outputs[0], skip_special_tokens=True)
        logger.info(f"Generated text: {result}")
        return {"generated_text": result}
    except Exception as e:
        logger.error(f"Generation error: {e}")
        return {"error": str(e)}