fix

app.py

@@ -1,29 +1,46 @@
-import gradio as gr
-from transformers import AutoProcessor, Gemma3ForConditionalGeneration
-import torch
 import os
+import torch
+from fastapi import FastAPI
+from transformers import AutoProcessor, AutoModelForCausalLM
+from pydantic import BaseModel
+import logging
 
-print(os.getenv("HF_TOKEN"))
+# ログ設定
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# モデルロード
-model_name = "unsloth/gemma-3-4b-it"
-processor = AutoProcessor.from_pretrained(model_name)
-model = Gemma3ForConditionalGeneration.from_pretrained(
-    model_name, torch_dtype=torch.bfloat16, device_map="auto"
-)
+app = FastAPI()
 
-def generate_text(text, max_length=50):
-    inputs = processor(text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-    outputs = model.generate(**inputs, max_length=max_length)
-    return processor.decode(outputs[0], skip_special_tokens=True)
+# モデルロード
+model_name = "google/gemma-3-4b-it"
+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
+    )
+    logger.info("Model loaded successfully")
+except Exception as e:
+    logger.error(f"Model load error: {e}")
+    raise
 
-# Gradioインターフェース
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=["text", "slider"],
-    outputs="text",
-    title="Gemma 3 API"
-)
+class TextInput(BaseModel):
+    text: str
+    max_length: int = 50
 
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+@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)}