Update app.py

app.py

@@ -1,88 +1,54 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel, PeftConfig
-from fastapi.middleware.cors import CORSMiddleware
 import torch
 import os
 
-app = FastAPI()
-
-# Allow CORS (customize in production)
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Hugging Face access token (from env)
+# Hugging Face access token (stored in HF Spaces secrets)
 hf_token = os.getenv("HF_TOKEN")
 
-# HF model repo (includes adapter + full model)
 adapter_path = "imnim/multi-label-email-classifier"
 
-try:
-    # Load PEFT adapter config
-    peft_config = PeftConfig.from_pretrained(adapter_path, token=hf_token)
-
-    # Try loading in bfloat16, fallback to float32
-    try:
-        base_model = AutoModelForCausalLM.from_pretrained(
-            peft_config.base_model_name_or_path,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            token=hf_token
-        )
-    except Exception:
-        base_model = AutoModelForCausalLM.from_pretrained(
-            peft_config.base_model_name_or_path,
-            torch_dtype=torch.float32,
-            device_map="auto",
-            token=hf_token
-        )
+# Load PEFT config
+peft_config = PeftConfig.from_pretrained(adapter_path, token=hf_token)
 
-    tokenizer = AutoTokenizer.from_pretrained(
+# Load base model (fallback to float32 if bfloat16 fails)
+try:
+    base_model = AutoModelForCausalLM.from_pretrained(
         peft_config.base_model_name_or_path,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
         token=hf_token
     )
-
-    # Load the adapter
-    model = PeftModel.from_pretrained(
-        base_model,
-        adapter_path,
+except:
+    base_model = AutoModelForCausalLM.from_pretrained(
+        peft_config.base_model_name_or_path,
+        torch_dtype=torch.float32,
+        device_map="auto",
         token=hf_token
     )
 
-    # Create the pipeline — no device argument (handled by accelerate)
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-except Exception as e:
-    raise RuntimeError(f"❌ Failed to load model + adapter: {str(e)}")
+tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path, token=hf_token)
 
-# === Request Schema ===
-class EmailInput(BaseModel):
-    subject: str
-    body: str
+model = PeftModel.from_pretrained(base_model, adapter_path, token=hf_token)
 
-# === Endpoint ===
-@app.post("/classify")
-async def classify_email(data: EmailInput):
-    prompt = f"""### Subject:\n{data.subject}\n\n### Body:\n{data.body}\n\n### Labels:"""
-    try:
-        result = pipe(prompt, max_new_tokens=50, do_sample=True, top_k=50, top_p=0.95)
-        full_text = result[0]["generated_text"]
-        label_section = full_text.split("### Labels:")[1].strip()
-        return {"label": label_section}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Model inference failed: {str(e)}")
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+# Define classification function
+def classify_email(subject, body):
+    prompt = f"""### Subject:\n{subject}\n\n### Body:\n{body}\n\n### Labels:"""
+    result = pipe(prompt, max_new_tokens=50, do_sample=True, top_k=50, top_p=0.95)
+    full_text = result[0]["generated_text"]
+    label_section = full_text.split("### Labels:")[1].strip()
+    return label_section
 
+# Gradio UI
+demo = gr.Interface(
+    fn=classify_email,
+    inputs=["text", "text"],
+    outputs="text",
+    title="Multi-label Email Classifier",
+    description="Enter subject and body to get label prediction"
+)
 
-
-import uvicorn
-
-if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")
-        
+demo.launch()