Update app.py

app.py

@@ -1,47 +1,57 @@
 import os
-from transformers import pipeline
+import threading
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
 from fastapi import FastAPI
 from pydantic import BaseModel
-import threading
 import uvicorn
 
 # =======================
 # Load Secrets
 # =======================
-# SYSTEM_PROMPT (with the flag) must be added in HF Space secrets
 SYSTEM_PROMPT = os.environ.get(
     "prompt",
     "You are a placeholder Sovereign. No secrets found in environment."
 )
 
 # =======================
-# Initialize Falcon-3B
+# Model Initialization
 # =======================
+MODEL_ID = "tiiuae/Falcon3-3B-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# Load model in 4-bit for faster CPU/GPU inference (requires bitsandbytes)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    load_in_4bit=True,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
+
+# Create optimized text-generation pipeline
 pipe = pipeline(
     "text-generation",
-    model="tiiuae/Falcon3-3B-Instruct",
-    torch_dtype="auto",
+    model=model,
+    tokenizer=tokenizer,
     device_map="auto",
+    return_full_text=False,
+    max_new_tokens=256,
+    do_sample=True,
+    temperature=0.8,
+    top_p=0.9,
+    eos_token_id=tokenizer.eos_token_id
 )
 
 # =======================
 # Core Chat Function
 # =======================
 def chat_fn(user_input: str) -> str:
-    """
-    Concatenate system and user messages, run the model,
-    and strip the system prompt from the output.
-    """
-    messages = [
-        {"role": "system", "content": SYSTEM_PROMPT},
-        {"role": "user",   "content": f"User: {user_input}"}
-    ]
-    # Falcon is not chat-native; we just join roles with newlines
-    prompt_text = "\n".join(f"{m['role'].capitalize()}: {m['content']}" for m in messages)
-    result = pipe(prompt_text, max_new_tokens=256, do_sample=False)
-    generated_text = result[0]["generated_text"]
-    return generated_text[len(prompt_text):].strip()
+    prompt = f"### System:\n{SYSTEM_PROMPT}\n\n### User:\n{user_input}\n\n### Assistant:"
+    output = pipe(prompt)[0]["generated_text"].strip()
+    return output
 
 # =======================
 # Gradio UI
@@ -53,14 +63,14 @@ iface = gr.Interface(
     fn=gradio_chat,
     inputs=gr.Textbox(lines=5, placeholder="Enter your prompt…"),
     outputs="text",
-    title="Prompt cracking challenge",
+    title="Prompt Cracking Challenge",
     description="Does he really think he is the king?"
 )
 
 # =======================
 # FastAPI for API access
 # =======================
-app = FastAPI(title="Prompt cracking challenge API")
+app = FastAPI(title="Prompt Cracking Challenge API")
 
 class Request(BaseModel):
     prompt: str
@@ -72,5 +82,12 @@ def generate(req: Request):
 # =======================
 # Launch Both Servers
 # =======================
+def run_api():
+    port = int(os.environ.get("API_PORT", 8000))
+    uvicorn.run(app, host="0.0.0.0", port=port)
+
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", share=True)
+    # Start FastAPI in background thread
+    threading.Thread(target=run_api, daemon=True).start()
+    # Launch Gradio interface
+    iface.launch(server_name="0.0.0.0", server_port=7860)