fix

app.py

@@ -1,12 +1,15 @@
 import os
 import torch
-import gradio as gr
+from fastapi import FastAPI
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from pydantic import BaseModel
 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}")
@@ -32,26 +35,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").to("cuda" if torch.cuda.is_available() else "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").to("cuda" if torch.cuda.is_available() else "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__":
-    logger.info("Launching Gradio interface")
-    iface.launch(server_name="0.0.0.0", server_port=8080, share=True)
-    logger.info("end")
+        return {"error": str(e)}

requirements.txt

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