Update app.py

app.py

@@ -1,11 +1,8 @@
-from fastapi import FastAPI
+import spaces
+import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
-from pydantic import BaseModel
-import uvicorn
-
-app = FastAPI()
 
 BASE_MODEL = "deepseek-ai/deepseek-coder-1.3b-instruct"
 ADAPTER = "praveends/migration-copilot-deepseek-coder-1-3b-instruct"
@@ -19,24 +16,16 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float32,
     trust_remote_code=True,
     low_cpu_mem_usage=True,
-    device_map="cpu",
 )
-model = PeftModel.from_pretrained(model, ADAPTER, device_map="cpu")
+model = PeftModel.from_pretrained(model, ADAPTER)
 model = model.merge_and_unload()
 model.eval()
 print("Model loaded!")
 
-class GenerateRequest(BaseModel):
-    prompt: str
-
-@app.get("/health")
-def health():
-    return {"status": "healthy"}
-
-@app.post("/generate")
-def generate(request: GenerateRequest):
+@spaces.GPU(duration=60)
+def generate(prompt: str) -> str:
     inputs = tokenizer(
-        request.prompt, return_tensors="pt",
+        prompt, return_tensors="pt",
         truncation=True, max_length=512
     )
     with torch.no_grad():
@@ -48,8 +37,14 @@ def generate(request: GenerateRequest):
             pad_token_id=tokenizer.eos_token_id,
         )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    generated = response[len(request.prompt):].strip()
-    return {"generated": generated}
+    return response[len(prompt):].strip()
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=gr.Textbox(label="Prompt", lines=10),
+    outputs=gr.Textbox(label="Generated PySpark", lines=10),
+    title="Migration Copilot Inference",
+    api_name="generate",
+)
 
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+demo.launch()