fix

app.py

@@ -1,47 +1,56 @@
 import os
 import torch
-from fastapi import FastAPI
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from pydantic import BaseModel
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import logging
 
-# ログ設定
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
-
-# モデルロード
 model_name = "google/gemma-2-2b-it"
 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",
         token=os.getenv("HF_TOKEN"),
         low_cpu_mem_usage=True,
-        load_in_4bit=True
+        quantization_config=quantization_config
     )
     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").to("cuda" if torch.cuda.is_available() else "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__":
+    logger.info("Launching Gradio interface")
+    iface.launch(server_name="0.0.0.0", server_port=8080)

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