Update app.py

app.py

@@ -1,78 +1,33 @@
-import os
 from fastapi import FastAPI, Request
-from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from fastapi.responses import JSONResponse
+from huggingface_hub import InferenceClient
 
-import os
+# Initialize FastAPI app
+app = FastAPI()
 
-# ======================
-# ✅ Set writable cache paths
-# ======================
-custom_cache = "/app/hf_cache"
-os.environ["HF_HOME"] = custom_cache
-os.environ["TRANSFORMERS_CACHE"] = os.path.join(custom_cache, "transformers")
-os.environ["HF_DATASETS_CACHE"] = os.path.join(custom_cache, "datasets")
-os.environ["HF_HUB_CACHE"] = os.path.join(custom_cache, "hub")
-
-os.makedirs(custom_cache, exist_ok=True)
-os.makedirs(os.environ["TRANSFORMERS_CACHE"], exist_ok=True)
-os.makedirs(os.environ["HF_DATASETS_CACHE"], exist_ok=True)
-os.makedirs(os.environ["HF_HUB_CACHE"], exist_ok=True)
-
-print("✅ Hugging Face cache directories set to:", custom_cache)
-
-# ============================
-# 🚀 Model Setup
-# ============================
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
-
-print(f"🔥 Loading {MODEL_NAME} ...")
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype="auto",
-    device_map="auto"
-)
-
-generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    temperature=0.7,
-    top_p=0.95
-)
-
-# ============================
-# 🌐 FastAPI Setup
-# ============================
-app = FastAPI(title="Jarvis AI V2", version="1.0")
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# Use the instruction-tuned Gemma model (adjust name if 3B or 9B version is available)
+MODEL_NAME = "google/gemma-2-9b-it"
+client = InferenceClient(model=MODEL_NAME)
 
 @app.get("/")
-def root():
-    return {"message": "🤖 Jarvis AI V2 is running successfully!"}
+async def home():
+    return {"status": "online", "model": MODEL_NAME}
 
-@app.post("/generate")
-async def generate(request: Request):
+@app.post("/chat")
+async def chat(request: Request):
     data = await request.json()
     prompt = data.get("prompt", "")
-    if not prompt:
-        return {"error": "Missing prompt text."}
-
-    print(f"🧠 Generating for prompt: {prompt[:50]}...")
-
-    output = generator(prompt)[0]["generated_text"]
-    return {"response": output.strip()}
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    if not prompt.strip():
+        return JSONResponse({"error": "Prompt is empty"}, status_code=400)
+
+    try:
+        response = client.text_generation(
+            prompt,
+            max_new_tokens=300,
+            temperature=0.7,
+            top_p=0.9
+        )
+        return JSONResponse({"response": response})
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)