Update app.py

app.py

@@ -1,19 +1,20 @@
+import os
+from pathlib import Path
 from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
-from pathlib import Path
-import os
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
-# Set Hugging Face cache dir
+# Setup Hugging Face cache directory
 cache_dir = os.getenv("TRANSFORMERS_CACHE", "/cache")
 os.makedirs(cache_dir, exist_ok=True)
 
-# Token for private models
+# Optional token (for private models)
 hf_token = os.getenv("HF_TOKEN")
 
-# Load model
+# Load model and tokenizer
 model_id = "mistralai/Mistral-7B-Instruct-v0.2"
+
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -23,17 +24,26 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype="auto"
 )
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
+# Build generation pipeline
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=256,
+    temperature=0.7,
+    top_p=0.9,
+    repetition_penalty=1.1
+)
 
-# Create FastAPI app
+# Initialize FastAPI app
 app = FastAPI()
 
-# Serve static files (JS, CSS, etc.)
+# Serve static files like script.js
 app.mount("/static", StaticFiles(directory="."), name="static")
 
 # Route: Serve index.html at root
 @app.get("/", response_class=HTMLResponse)
-async def root():
+async def serve_home():
     html_path = Path("index.html")
     return HTMLResponse(content=html_path.read_text(), status_code=200)
 
@@ -47,5 +57,13 @@ async def ask_ai(request: Request):
         return JSONResponse(content={"answer": "❗ Please enter a valid question."})
 
     prompt = f"[INST] {question} [/INST]"
-    output = pipe(prompt)[0]["generated_text"]
-    return JSONResponse(content={"answer": output.strip()})
+    try:
+        output = pipe(prompt)[0]["generated_text"]
+        return JSONResponse(content={"answer": output.strip()})
+    except Exception as e:
+        return JSONResponse(content={"answer": f"⚠️ Error: {str(e)}"})
+
+# Optional: Serve script.js if not using /static path in HTML
+@app.get("/script.js")
+async def serve_script():
+    return FileResponse("script.js", media_type="application/javascript")