fix

app.py

@@ -1,7 +1,8 @@
 import os
 import torch
-import gradio as gr
+from fastapi import FastAPI
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from pydantic import BaseModel
 import logging
 
 logging.basicConfig(level=logging.INFO)
@@ -16,8 +17,6 @@ model = None
 try:
     logger.info(f"Loading model: {model_name}")
     tokenizer = AutoTokenizer.from_pretrained(model_name, token=os.getenv("HF_TOKEN"))
-    use_gpu = torch.cuda.is_available()
-    logger.info(f"GPU available: {use_gpu}")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float16,  # メモリ削減
@@ -30,29 +29,19 @@ except Exception as e:
     logger.error(f"Model load error: {e}")
     raise
 
-def generate_text(text, max_length=50):
+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: {text}")
-        inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True).to("cpu")
-        outputs = model.generate(**inputs, max_length=max_length)
+        logger.info(f"Generating text for input: {input.text}")
+        inputs = tokenizer(input.text, return_tensors="pt", max_length=512, truncation=True).to("cpu")
+        outputs = model.generate(**inputs, max_length=input.max_length)
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         logger.info(f"Generated text: {result}")
-        return result
+        return {"generated_text": result}
     except Exception as e:
         logger.error(f"Generation error: {e}")
-        return f"Error: {str(e)}"
-
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=[gr.Textbox(label="Input Text"), gr.Slider(10, 100, value=50, label="Max Length")],
-    outputs=gr.Textbox(label="Generated Text"),
-    title="Gemma 2 API"
-)
-
-if __name__ == "__main__":
-    try:
-        logger.info("Launching Gradio interface")
-        iface.launch(server_name="0.0.0.0", server_port=8080)
-    except Exception as e:
-        logger.error(f"Gradio launch error: {e}")
-        raise
+        return {"error": str(e)}"