Update app.py

app.py

@@ -10,7 +10,7 @@ import re  # for parsing <think> blocks
 import gradio as gr
 import torch
 from transformers import pipeline, TextIteratorStreamer
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from duckduckgo_search import DDGS
 
 from transformers import modeling_utils
@@ -55,53 +55,34 @@ def load_pipeline(model_name):
         return PIPELINES[model_name]
     repo = MODELS[model_name]["repo_id"]
     if model_name == "secgpt-mini":
-        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True, subfolder="models")
+        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True, device_map=device, subfolder="models")
+        model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, device_map=device, subfolder="models")
     else:
-        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True, device_map=device)
+        model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, device_map=device)
     for dtype in (torch.bfloat16, torch.float16, torch.float32):
         try:
-            if model_name == "secgpt-mini":
-                pipe = pipeline(
+            pipe = pipeline(
                     task="text-generation",
-                    model=repo,
+                    model=model,
                     tokenizer=tokenizer,
                     trust_remote_code=True,
                     torch_dtype=dtype,
                     device_map=device,
-                    subfolder="models"
-                )
-            else:
-                pipe = pipeline(
-                    task="text-generation",
-                    model=repo,
-                    tokenizer=tokenizer,
-                    trust_remote_code=True,
-                    torch_dtype=device,
-                    device_map="auto",
                 )
             PIPELINES[model_name] = pipe
             return pipe
         except Exception:
             continue
     # Final fallback
-    if model_name == "secgpt-mini":
-        pipe = pipeline(
-                task="text-generation",
-                model=repo,
-                tokenizer=tokenizer,
-                trust_remote_code=True,
-                torch_dtype=dtype,
-                device_map=device,
-                subfolder="models"
-        )
-    else:
-        pipe = pipeline(
+    pipe = pipeline(
             task="text-generation",
-            model=repo,
+            model=model,
             tokenizer=tokenizer,
             trust_remote_code=True,
-            device_map=device
-        )
+            torch_dtype=dtype,
+            device_map=device,
+    )
     PIPELINES[model_name] = pipe
     return pipe
 
@@ -290,9 +271,9 @@ with gr.Blocks(title="Yee R1 Demo") as demo:
             gr.Markdown("### Generation Parameters")
             max_tok = gr.Slider(64, 16384, value=4096, step=32, label="Max Tokens")
             temp = gr.Slider(0.1, 2.0, value=0.6, step=0.1, label="Temperature")
-            k = gr.Slider(1, 100, value=40, step=1, label="Top-K")
-            p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
-            rp = gr.Slider(1.0, 2.0, value=1.2, step=0.1, label="Repetition Penalty")
+            k = gr.Slider(1, 100, value=20, step=1, label="Top-K")
+            p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P")
+            rp = gr.Slider(1.0, 2.0, value=1.0, step=0.1, label="Repetition Penalty")
             gr.Markdown("### Web Search Settings")
             mr = gr.Number(value=6, precision=0, label="Max Results")
             mc = gr.Number(value=600, precision=0, label="Max Chars/Result")