Update app.py

app.py

@@ -2,118 +2,101 @@ import gradio as gr
 import subprocess
 from huggingface_hub import hf_hub_download
 
-# 1. Install llama-cpp-python in runtime
+# --- 1. Setup & Install ---
 subprocess.run("pip install -q 'llama_cpp_python==0.3.15'", shell=True, check=False)
-
 from llama_cpp import Llama
 
-# 2. Load GGUF model from Hugging Face
+# --- 2. Load Model (GGUF) ---
 MODEL_REPO = "Jeppcode/ScalableLab2"
 GGUF_FILENAME = "model-q4_k_m.gguf"
 
 print(f"Downloading GGUF model {MODEL_REPO}/{GGUF_FILENAME} ...")
-model_path = hf_hub_download(
-    repo_id=MODEL_REPO,
-    filename=GGUF_FILENAME,
-)
+try:
+    model_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=GGUF_FILENAME,
+    )
+except Exception as e:
+    print(f"Error downloading model: {e}")
+    model_path = ""
 
-print("Initializing llama.cpp LLM ...")
-llm = Llama(
-    model_path=model_path,
-    n_ctx=2048,
-    n_threads=2,
-    n_batch=64,
-    use_mmap=True,
-    use_mlock=False,
-)
+llm = None
+if model_path:
+    print("Initializing llama.cpp LLM ...")
+    llm = Llama(
+        model_path=model_path,
+        n_ctx=2048,
+        n_threads=2,
+        n_batch=64,
+        use_mmap=True,
+        use_mlock=False,
+    )
 
-# Hardcoded system prompt (since style selector is removed)
-SYSTEM_PROMPT = "You are a helpful teaching assistant. Give clear, structured and detailed explanations."
+# --- 3. Style / System Prompts ---
+# These are the "Buttons" logic to change how the AI behaves
+STYLE_SYSTEM_PROMPTS = {
+    "Default": "You are a helpful, polite assistant.",
+    "Short answer": "Answer as concisely as possible, usually in 1–3 sentences.",
+    "Detailed explanation": "Give clear, structured and detailed explanations.",
+    "Step-by-step reasoning": "Think step by step and explain your reasoning clearly.",
+}
 
-def _extract_text_from_content(content):
-    """
-    Extracts text from Gradio 6 message format.
-    """
+def _extract_text(content):
     if isinstance(content, list):
-        texts = []
-        for block in content:
-            if isinstance(block, dict) and block.get("type") == "text":
-                texts.append(block.get("text", ""))
-            else:
-                texts.append(str(block))
-        return "\n".join(t for t in texts if t)
-    else:
-        return str(content)
+        return "\n".join(b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text")
+    return str(content)
 
-def build_prompt(message, history):
-    """
-    Builds the prompt using the hardcoded system prompt.
-    """
-    prompt_parts = []
-    prompt_parts.append(f"System: {SYSTEM_PROMPT}\n")
-    prompt_parts.append("Conversation:\n")
+def chat_fn(message, history, max_new_tokens, style):
+    if not llm: return "Error: Model not loaded."
+    
+    # Select the specific system prompt based on the button chosen
+    system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])
     
+    prompt = f"System: {system_prompt}\nConversation:\n"
     for msg in history or []:
         role = msg.get("role")
-        content = _extract_text_from_content(msg.get("content", ""))
-        if not content:
-            continue
-        if role == "user":
-            prompt_parts.append(f"User: {content}\n")
-        elif role == "assistant":
-            prompt_parts.append(f"Assistant: {content}\n")
-        elif role == "system":
-            prompt_parts.append(f"System (previous): {content}\n")
+        txt = _extract_text(msg.get("content", ""))
+        if txt:
+            if role == "user": prompt += f"User: {txt}\n"
+            elif role == "assistant": prompt += f"Assistant: {txt}\n"
             
-    # Current user message
-    prompt_parts.append(f"User: {message}\n")
-    prompt_parts.append("Assistant:")
-    full_prompt = "".join(prompt_parts)
-    return full_prompt
-
-def chat_fn(message, history, max_new_tokens):
-    """
-    Main chat function. 
-    Only accepts max_new_tokens as an additional input now.
-    """
-    prompt = build_prompt(message, history)
-    
-    # Internal defaults for the removed sliders
-    temperature = 0.7
-    top_p = 0.9
-    repetition_penalty = 1.0
+    prompt += f"User: {message}\nAssistant:"
 
+    # Default internal values for randomness
     output = llm(
         prompt,
         max_tokens=int(max_new_tokens),
-        temperature=temperature,
-        top_p=top_p,
-        repeat_penalty=repetition_penalty,
-        stop=["User:", "Assistant:", "System:", "Conversation:"],
+        temperature=0.7,
+        top_p=0.9,
+        stop=["User:", "Assistant:", "System:"],
     )
-    reply = output["choices"][0]["text"].strip()
-    return reply
+    return output["choices"][0]["text"].strip()
+
+# --- 4. UI Controls ---
 
-# 4. Only keep the Max Token slider
+# Slider for length
 max_new_tokens_slider = gr.Slider(
-    minimum=16,
-    maximum=256,
-    value=64,
-    step=8,
-    label="Max new tokens (response length)",
+    minimum=16, 
+    maximum=256, 
+    value=64, 
+    step=8, 
+    label="Max Response Length"
+)
+
+# The "Buttons" at the bottom for Style
+style_radio = gr.Radio(
+    choices=["Default", "Short answer", "Detailed explanation", "Step-by-step reasoning"],
+    value="Detailed explanation",
+    label="Answer Style"
 )
 
+# --- 5. Launch App (Clean / No Theme) ---
 demo = gr.ChatInterface(
     fn=chat_fn,
     title="Lab 2 – Fine-tuned GGUF model",
-    description=(
-        "Chat with our fine-tuned Llama-based model, converted to GGUF and "
-        "loaded via llama.cpp from Jeppcode/ScalableLab2."
-    ),
-    additional_inputs=[
-        max_new_tokens_slider,
-    ],
-    additional_inputs_accordion="Generation controls",
+    description="Chat with the fine-tuned Llama model. Use the controls below to change the response style.",
+    additional_inputs=[max_new_tokens_slider, style_radio],
+    additional_inputs_accordion="Controls",
 )
 
 if __name__ == "__main__":