Update app.py

app.py

@@ -2,80 +2,60 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ---------------------------------------------------------
-# 1) Points to your Hugging Face repo and subfolder:
-#    "wuhp/myr1" is the repository
-#    "myr1" is the subfolder where the config/tokenizer/model are located.
-# ---------------------------------------------------------
 MODEL_REPO = "wuhp/myr1"
 SUBFOLDER = "myr1"
 
-# ---------------------------------------------------------
-# 2) Load the tokenizer and model from the Hub
-#    - trust_remote_code=True allows custom config & modeling files.
-# ---------------------------------------------------------
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_REPO,
-    subfolder=SUBFOLDER,         # important because the model files sit inside 'myr1'
+    subfolder=SUBFOLDER,
     trust_remote_code=True
 )
 
+# If your GPU has <24GB VRAM, consider 8-bit or CPU offloading
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_REPO,
-    subfolder=SUBFOLDER,         # also needed here
+    subfolder=SUBFOLDER,
     trust_remote_code=True,
-    device_map="auto",           # automatically place model layers on GPU(s) if available
-    torch_dtype=torch.float16,    # or "auto", "float32", "bfloat16", etc. as your hardware supports
+    device_map="auto",           # tries to place layers on GPU, then CPU if needed
+    torch_dtype=torch.float16,    # or bfloat16 or float32
     low_cpu_mem_usage=True
 )
 
-# Put the model in evaluation mode
 model.eval()
 
-
-def generate_text(prompt, max_length=128, temperature=0.7, top_p=0.9):
-    """
-    Generate text from your DeepSeekR1 model, given an input prompt.
-    """
-    # Convert to token IDs and move to model device (GPU/CPU)
+def generate_text(prompt, max_length=64, temperature=0.7, top_p=0.9):
+    print("=== Starting generation ===")
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    # Generate output
-    with torch.no_grad():
+    try:
         output_ids = model.generate(
             **inputs,
-            max_length=max_length,
+            max_new_tokens=max_length,   # alternative to max_length
             temperature=temperature,
             top_p=top_p,
             do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,  # or set to a real pad_token_id if your model uses one
+            pad_token_id=tokenizer.eos_token_id
         )
-
-    # Decode the tokens back into a string
+        print("=== Generation complete ===")
+    except Exception as e:
+        print(f"Error during generation: {e}")
+        return str(e)
     return tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-
-# ---------------------------------------------------------
-# 3) Build Gradio UI
-# ---------------------------------------------------------
 demo = gr.Interface(
     fn=generate_text,
     inputs=[
         gr.Textbox(
-            lines=5,
-            label="Enter your prompt",
-            placeholder="Type something for the DeepSeek model..."
+            lines=4,
+            label="Prompt",
+            placeholder="Try a short prompt, e.g., Hello!"
         ),
-        gr.Slider(64, 1024, step=1, value=128, label="Max Length"),
-        gr.Slider(0.0, 1.5, step=0.1, value=0.7, label="Temperature"),
-        gr.Slider(0.0, 1.0, step=0.05, value=0.9, label="Top-p"),
+        gr.Slider(8, 512, value=64, step=1, label="Max New Tokens"),
+        gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p"),
     ],
     outputs="text",
-    title="DeepSeek-R1 Gradio Demo",
-    description=(
-        "This Gradio interface loads the DeepSeek model from Hugging Face and lets you "
-        "generate text by entering a prompt. Adjust parameters to see how output changes."
-    )
+    title="DeepSeek R1 Demo",
+    description="Generates text using the large DeepSeek model."
 )
 
 if __name__ == "__main__":