Update app.py

app.py

@@ -3,81 +3,52 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
-# ===== MODEL LOAD=====
 BASE_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-LORA_PATH = "./deepseek-lab-assistant"
+LORA_PATH = "./"  
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-tokenizer = AutoTokenizer.from_pretrained(
-    BASE_MODEL,
-    trust_remote_code=True
-)
+#Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 tokenizer.pad_token = tokenizer.eos_token
 
+#Base model
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True
+    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+    device_map="auto"
 )
 
+
 model = PeftModel.from_pretrained(model, LORA_PATH)
 model.eval()
 
 
-# ===== CHAT FUNCTION =====
-def respond(
-    message,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    # history = [{"role": "user"/"assistant", "content": "..."}]
-
-    prompt = system_message + "\n\n"
-
-    for h in history:
-        prompt += f"{h['role'].capitalize()}: {h['content']}\n"
-
-    prompt += f"User: {message}\nAssistant:"
-
+def chat(prompt, max_new_tokens=256, temperature=0.7):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
         output = model.generate(
             **inputs,
-            max_new_tokens=max_tokens,
+            max_new_tokens=max_new_tokens,
             temperature=temperature,
-            top_p=top_p,
-            do_sample=temperature > 0,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id
         )
 
-    text = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    if "Assistant:" in text:
-        text = text.split("Assistant:")[-1].strip()
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
-    return text
 
-
-# ===== GRADIO UI =====
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(
-            value="You are a helpful lab assistant. Explain ideas clearly. Do not rush to final answers.",
-            label="System message",
-        ),
-        gr.Slider(1, 1024, value=256, step=1, label="Max new tokens"),
-        gr.Slider(0.0, 1.5, value=0.3, step=0.05, label="Temperature"),
-        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
+# Gradio UI
+demo = gr.Interface(
+    fn=chat,
+    inputs=[
+        gr.Textbox(lines=5, label="Prompt"),
+        gr.Slider(1, 1024, value=256, label="Max tokens"),
+        gr.Slider(0.1, 1.5, value=0.7, label="Temperature"),
     ],
+    outputs=gr.Textbox(lines=10, label="Output"),
+    title="DeepSeek Lab Assistant (LoRA)",
 )
 
-with gr.Blocks() as demo:
-    chatbot.render()
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(True)