Update app.py

app.py

@@ -1,57 +1,103 @@
 
-from flask import Flask, request, jsonify
+import gradio as gr
 import logging
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import os
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+import uvicorn
 
-app = Flask(__name__)
+# Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
+# Load Hugging Face token from environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    logging.warning("HF_TOKEN not set. Private repositories may be inaccessible.")
+
+# Model configuration
+MODEL_PATH = "your_username/fine_tuned_bart_construction"  # Replace with your model path
+
 try:
-    # Load fine-tuned model and tokenizer
-    model = AutoModelForSeq2SeqLM.from_pretrained("./fine_tuned_model")
-    tokenizer = AutoTokenizer.from_pretrained("./fine_tuned_model")
-    logging.info("Model and tokenizer loaded successfully")
+    # Load model and tokenizer with authentication
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, use_auth_token=HF_TOKEN)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_auth_token=HF_TOKEN)
+    logging.info(f"Model and tokenizer loaded successfully from {MODEL_PATH}")
 except Exception as e:
     logging.error(f"Failed to load model or tokenizer: {e}")
     raise e
 
-@app.route("/generate", methods=["POST"])
-def generate_checklist():
+# Define input model for FastAPI
+class ChecklistInput(BaseModel):
+    role: str = "Supervisor"
+    project_id: str = "Unknown"
+    project_name: str = "Unknown Project"
+    milestones: str = "No milestones provided"
+
+# Initialize FastAPI for JSON endpoint (Gradio mounts it)
+app = FastAPI()
+
+@app.post("/generate")
+async def generate_checklist(data: ChecklistInput):
     try:
-        data = request.get_json()
-        if not data:
-            logging.warning("No JSON data provided in request")
-            return jsonify({"error": "Invalid or missing JSON data"}), 400
-        
-        role = data.get("role", "Supervisor")
-        project_id = data.get("project_id", "Unknown")
-        project_name = data.get("project_name", "Unknown Project")
-        milestones = data.get("milestones", "No milestones provided")
-        
         # Prepare input for the model
-        inputs = f"Role: {role} Project: {project_id} ({project_name}) Milestones: {milestones}"
+        inputs = f"Role: {data.role} Project: {data.project_id} ({data.project_name}) Milestones: {data.milestones}"
         logging.info(f"Generating checklist for inputs: {inputs[:100]}...")  # Truncate for logging
-        
+
         # Tokenize and generate
         input_ids = tokenizer(inputs, return_tensors="pt", max_length=128, truncation=True).input_ids
         outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True)
         checklist = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Static tips (can be enhanced with model output)
+
+        # Static tips (can be enhanced with model)
         tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress"
-        
+
         logging.info("Checklist and tips generated successfully")
-        return jsonify({
+        return {
             "checklist": checklist,
             "tips": tips
-        }), 200
+        }
     except Exception as e:
         logging.error(f"Error generating checklist: {e}")
-        return jsonify({"error": str(e)}), 500
+        return {"error": str(e)}
+
+# Gradio interface function
+def gradio_generate_checklist(role, project_id, project_name, milestones):
+    try:
+        inputs = f"Role: {role} Project: {project_id} ({project_name}) Milestones: {milestones}"
+        input_ids = tokenizer(inputs, return_tensors="pt", max_length=128, truncation=True).input_ids
+        outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True)
+        checklist = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress"
+        return checklist, tips
+    except Exception as e:
+        return f"Error: {str(e)}", ""
+
+# Define Gradio interface
+iface = gr.Interface(
+    fn=gradio_generate_checklist,
+    inputs=[
+        gr.Textbox(label="Role", value="Supervisor"),
+        gr.Textbox(label="Project ID", value="P001"),
+        gr.Textbox(label="Project Name", value="Building A"),
+        gr.Textbox(label="Milestones", value="Complete foundation by 5/15")
+    ],
+    outputs=[
+        gr.Textbox(label="Checklist"),
+        gr.Textbox(label="Tips")
+    ],
+    title="AI Coach for Site Supervisors",
+    description="Generate daily checklists and tips for site supervisors."
+)
+
+# Mount FastAPI to Gradio
+iface.app = app
 
 if __name__ == "__main__":
     try:
-        app.run(host="0.0.0.0", port=8000, debug=False)
+        # Launch Gradio with public URL for Hugging Face Spaces
+        iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+        logging.info("Gradio interface launched successfully")
     except Exception as e:
-        logging.error(f"Failed to start Flask server: {e}")
+        logging.error(f"Failed to launch Gradio interface: {e}")
         raise e