Update app.py

app.py

@@ -8,6 +8,7 @@ from peft import PeftModel
 from huggingface_hub import login, hf_hub_download
 from fastapi.responses import Response
 import sys
+import fastapi
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -30,14 +31,19 @@ def load_base_model():
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
-    base = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID, trust_remote_code=True, device_map="cpu", torch_dtype=torch.float32, token=HF_TOKEN)
+    base = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL_ID,
+        trust_remote_code=True,
+        device_map="cpu",
+        torch_dtype=torch.float32,
+        token=HF_TOKEN
+    )
     return tokenizer, base
 
 def load_lora_model(base_model):
     cfg_path = hf_hub_download(repo_id=LORA_REPO, filename="adapter_config.json", token=HF_TOKEN)
     with open(cfg_path) as f:
         cfg = json.load(f)
-    # Clean any extra keys
     cfg.pop("layer_replication", None)
     cfg.pop("alpha_pattern", None)
     model = PeftModel.from_pretrained(base_model, LORA_REPO, token=HF_TOKEN)
@@ -55,24 +61,31 @@ def load_model():
 tokenizer, model = load_model()
 
 def generate_response(prompt):
-    safe = prompt.replace('"','').replace("\\","")
+    safe = prompt.replace('"', '').replace("\\", "")
     text = f"""<|im_start|>system\nYou are a cybersecurity expert.<|im_end|>\n<|im_start|>user\n{safe}<|im_end|>\n<|im_start|>assistant\n"""
     inputs = tokenizer(text, return_tensors="pt").to("cpu")
     with torch.no_grad():
-        out = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.9, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id)
+        out = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id
+        )
     decoded = tokenizer.decode(out[0], skip_special_tokens=False)
     if "<|im_start|>assistant" in decoded:
         return decoded.split("<|im_start|>assistant\n")[1].split("<|im_end|>")[0].strip()
     return decoded
 
-# Use custom FastAPI app with proper response
-import fastapi
 app = fastapi.FastAPI()
 
 @app.get("/ping")
 def ping():
     return Response(status_code=204)
 
+# Use correct path prefix to avoid 404 for static files
 gr_app = gr.Interface(
     fn=generate_response,
     inputs=gr.Textbox(lines=5, label="Prompt"),