fix

app.py

@@ -1,8 +1,7 @@
 import os
 import torch
-from fastapi import FastAPI
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-from pydantic import BaseModel
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 
 logging.basicConfig(level=logging.INFO)
@@ -11,43 +10,49 @@ logger = logging.getLogger(__name__)
 app = FastAPI()
 
 model_name = "google/gemma-2-2b-it"
+tokenizer = None
+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}")
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.bfloat16,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_use_double_quant=True
-    ) if use_gpu else None
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
+        torch_dtype=torch.float16,  # メモリ削減
+        device_map="cpu",  # GPU利用不可
         token=os.getenv("HF_TOKEN"),
-        low_cpu_mem_usage=True,
-        quantization_config=quantization_config
+        low_cpu_mem_usage=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):
+def generate_text(text, max_length=50):
     try:
-        logger.info(f"Generating text for input: {input.text}")
-        inputs = tokenizer(input.text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-        outputs = model.generate(**inputs, max_length=input.max_length)
+        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)
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         logger.info(f"Generated text: {result}")
-        return {"generated_text": result}
+        return result
     except Exception as e:
         logger.error(f"Generation error: {e}")
-        return {"error": str(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

requirements.txt

@@ -6,3 +6,4 @@ bitsandbytes==0.42.0
 accelerate==0.26.1
 fastapi==0.115.0
 uvicorn==0.30.6
+gradio==4.15.0