Update app.py

app.py

@@ -8,19 +8,29 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 MODEL_ID = "DavidBazaldua/llama3_finetuned_transformes"
 
-DEVICE = "cpu"  # Space on CPU
-DTYPE = torch.float32  # safer on CPU
+DEVICE = "cpu"          # Space on CPU
+DTYPE = torch.float32   # safer on CPU
 
+# Limit CPU threads (you can try 1, 2, 4, etc. depending on performance)
+torch.set_num_threads(2)
+
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=DTYPE,
-    low_cpu_mem_usage=True,  # intenta optimizar memoria
+    low_cpu_mem_usage=True,  # optimize memory usage on CPU
 )
 model.to(DEVICE)
-model.eval()  # modo evaluación
-torch.set_num_threads(2)  # limitar número de threads (ajústalo si quieres)
+model.eval()
+
+DEFAULT_SYSTEM_PROMPT = (
+    "You are a helpful, knowledgeable, and professional AI assistant. "
+    "You respond in English unless the user explicitly requests another language. "
+    "Provide clear, concise answers and reason step by step when it is useful."
+)
+
 # ---------------------------------------------------------------------
 # Prompt building (using the chat template from the tokenizer)
 # ---------------------------------------------------------------------
@@ -75,6 +85,9 @@ def generate_answer(system_prompt, context, message, history, max_tokens, temper
     if not system_prompt or system_prompt.strip() == "":
         system_prompt = DEFAULT_SYSTEM_PROMPT
 
+    # Hard cap for safety on CPU
+    max_tokens = int(min(max_tokens, 128))
+
     prompt = build_prompt(system_prompt, context, history, message)
 
     inputs = tokenizer(
@@ -86,21 +99,22 @@ def generate_answer(system_prompt, context, message, history, max_tokens, temper
     with torch.no_grad():
         output_tokens = model.generate(
             **inputs,
-            max_new_tokens=int(max_tokens),
+            max_new_tokens=max_tokens,
             do_sample=True,
             temperature=float(temperature),
             top_p=float(top_p),
             pad_token_id=tokenizer.eos_token_id,
         )
 
-    generated = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-    # Since we used the chat template, the decode already gives us the full conversation
-    # but we only need the new assistant segment. For simplicity, take everything after the last user message.
-    answer = generated[len(prompt):].strip()
+    # Decode full text (prompt + completion)
+    full_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
 
-    if not answer:
-        # Fallback if slicing fails for any reason
-        answer = generated.strip()
+    # Try to extract only the assistant's new part
+    if full_text.startswith(prompt):
+        answer = full_text[len(prompt):].strip()
+    else:
+        # Fallback if for some reason the decoded text does not start with prompt
+        answer = full_text.strip()
 
     history = history + [[message, answer]]
     return answer, history
@@ -120,6 +134,7 @@ def chat(message, history, system_prompt, context, max_tokens, temperature, top_
         top_p=top_p,
     )
 
+    # Return empty input and updated history for the Chatbot
     return "", updated_history
 
 
@@ -164,10 +179,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
             max_tokens_slider = gr.Slider(
                 label="Max new tokens",
-                minimum=64,
-                maximum=512,
-                value=256,
-                step=32,
+                minimum=32,
+                maximum=256,
+                value=128,
+                step=16,
             )
             temperature_slider = gr.Slider(
                 label="Temperature",
@@ -218,3 +233,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
 if __name__ == "__main__":
     demo.launch()
+