Update app.py

app.py

@@ -1,61 +1,83 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer # Removed BitsAndBytesConfig as we are not quantizing for CPU
 from peft import PeftModel
 import torch
+import os # Ensure os is imported for potential path joining if needed
 
 # --- Configuration ---
 base_model_id = "Qwen/Qwen-1_8B-Chat"
 lora_adapter_id = "jinv2/qwen-1_8b-hemiplegia-lora" # Your HF Model ID
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# device = "cuda" if torch.cuda.is_available() else "cpu" # Will always be "cpu" on free tier
+device = "cpu" # Explicitly set to CPU for this configuration
 print(f"Using device: {device}")
 
 # --- Load Model and Tokenizer ---
 print("Loading tokenizer...")
 try:
+    # Try loading tokenizer from your LoRA repo first, as it might contain specific settings
     tokenizer = AutoTokenizer.from_pretrained(lora_adapter_id, trust_remote_code=True)
     print(f"Successfully loaded tokenizer from {lora_adapter_id}.")
-except Exception:
-    print(f"Could not load tokenizer from {lora_adapter_id}, falling back to {base_model_id}.")
+except Exception as e_lora_tok:
+    print(f"Could not load tokenizer from {lora_adapter_id} (Error: {e_lora_tok}), falling back to {base_model_id}.")
     tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)
 
+# Set pad_token if not already set
 if tokenizer.pad_token_id is None:
     if tokenizer.eos_token_id is not None:
         tokenizer.pad_token_id = tokenizer.eos_token_id
-    else: # Fallback for Qwen, ensure this ID is correct for your Qwen version
-        tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<|endoftext|>") if "<|endoftext|>" in tokenizer.vocab else 0
+        tokenizer.pad_token = tokenizer.eos_token
+        print(f"Set tokenizer.pad_token_id to eos_token_id: {tokenizer.pad_token_id}")
+    else:
+        # Fallback for Qwen, ensure this ID is correct for your Qwen version
+        try:
+            qwen_eos_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
+            tokenizer.pad_token_id = qwen_eos_id
+            tokenizer.pad_token = "<|endoftext|>"
+            print(f"Set tokenizer.pad_token_id to ID of '<|endoftext|>: {tokenizer.pad_token_id}")
+        except KeyError:
+            tokenizer.pad_token_id = 0 # Absolute fallback, very risky
+            tokenizer.pad_token = tokenizer.decode([0])
+            print(f"CRITICAL WARNING: Could not set pad_token_id reliably. Set to 0 ('{tokenizer.pad_token}').")
 
 tokenizer.padding_side = "left" # Important for generation
 
-print("Loading base model with quantization...")
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16 # As used in fine-tuning
-)
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,
-    quantization_config=quantization_config,
-    trust_remote_code=True,
-    device_map={"":0} if device == "cuda" else "cpu" # Load directly to GPU if available, else CPU
-)
-print("Base model loaded.")
+print("Loading base model (NO QUANTIZATION as running on CPU)...")
+# IMPORTANT: For CPU, we cannot use bitsandbytes 4-bit quantization.
+# We load the model in its original precision (or try float16/bfloat16 if memory allows and CPU supports).
+# This will be much slower and more memory-intensive.
+try:
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model_id,
+        trust_remote_code=True,
+        torch_dtype=torch.float32,  # Use float32 for CPU for max compatibility, bfloat16 might work on some newer CPUs
+        # device_map="auto" will likely map to CPU. Can be explicit: device_map="cpu"
+        device_map={"":device} # Ensure model parts are on the correct device
+    )
+    print("Base model loaded.")
+except Exception as e_load_model:
+    print(f"Error loading base model: {e_load_model}")
+    raise # Re-raise the exception to stop the app if model loading fails
 
 print(f"Loading LoRA adapter: {lora_adapter_id}...")
-model = PeftModel.from_pretrained(base_model, lora_adapter_id)
-model.eval() # Set to evaluation mode
-print("LoRA adapter loaded and model is ready.")
-if device == "cpu": # If on CPU, PEFT might not automatically move the full model if device_map wasn't used correctly for CPU
-    model = model.to(device)
-    print(f"Model explicitly moved to {device}")
+try:
+    # For CPU, PEFT should still work. The model should be on the CPU before applying adapter.
+    model = PeftModel.from_pretrained(base_model, lora_adapter_id)
+    model.eval() # Set to evaluation mode
+    model = model.to(device) # Ensure the final PEFT model is on the CPU
+    print("LoRA adapter loaded and model is on CPU, ready for inference.")
+except Exception as e_load_adapter:
+    print(f"Error loading LoRA adapter: {e_load_adapter}")
+    raise
 
 
 # --- Prediction Function ---
 def get_response(user_query):
     system_prompt_content = "你是一个专注于偏瘫、脑血栓、半身不遂领域的医疗问答助手。"
     
-    # Construct prompt using Qwen's ChatML format
     prompt = f"<|im_start|>system\n{system_prompt_content}<|im_end|>\n<|im_start|>user\n{user_query}<|im_end|>\n<|im_start|>assistant\n"
     
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512-150).to(model.device) # Leave space for generation
+    # Ensure inputs are on the same device as the model
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512-150).to(model.device)
 
     eos_token_ids_list = []
     if isinstance(tokenizer.eos_token_id, int):
@@ -65,25 +87,31 @@ def get_response(user_query):
         if im_end_token_id not in eos_token_ids_list:
             eos_token_ids_list.append(im_end_token_id)
     except KeyError: pass
-    if not eos_token_ids_list and tokenizer.eos_token_id is not None: # Fallback if list is empty but single eos_token_id exists
-         eos_token_ids_list = [tokenizer.eos_token_id]
-    elif not eos_token_ids_list: # Absolute fallback
-        print("Warning: EOS token ID list is empty. Generation might not stop correctly.")
-        # Attempt to use a known Qwen EOS ID if possible, otherwise generation might be problematic.
-        # This scenario should ideally be avoided by robust tokenizer setup.
-        # eos_token_ids_list = [tokenizer.vocab_size - 1] # Very risky fallback
+    
+    # Fallback if eos_token_ids_list is still empty
+    if not eos_token_ids_list:
+        if tokenizer.eos_token_id is not None:
+             eos_token_ids_list = [tokenizer.eos_token_id]
+        else:
+            print("Warning: EOS token ID list is empty and eos_token_id is None. Generation might not stop correctly.")
+            # Attempt to use a known Qwen EOS ID if possible, otherwise generation might be problematic.
+            try:
+                eos_token_ids_list = [tokenizer.convert_tokens_to_ids("<|endoftext|>")]
+            except KeyError:
+                eos_token_ids_list = [tokenizer.vocab_size - 1 if tokenizer.vocab_size else 0] # Very risky fallback
+
 
-    print(f"Generating response for query: '{user_query}'")
-    with torch.no_grad():
+    print(f"Generating response for query: '{user_query}' on device: {model.device}")
+    with torch.no_grad(): # Inference doesn't need gradient calculation
         outputs = model.generate(
             **inputs,
             max_new_tokens=150,
             pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=eos_token_ids_list if eos_token_ids_list else None, # Pass list or None
+            eos_token_id=eos_token_ids_list if eos_token_ids_list else None,
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
-            num_beams=1 # Use 1 for sampling, or >1 for beam search (do_sample=False then)
+            num_beams=1
         )
     
     response_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
@@ -95,15 +123,18 @@ iface = gr.Interface(
     fn=get_response,
     inputs=gr.Textbox(lines=3, placeholder="请输入您关于偏瘫、脑血栓或半身不遂的问题...", label="您的问题 (Your Question)"),
     outputs=gr.Textbox(lines=5, label="模型回答 (Model Response)"),
-    title="偏瘫脑血栓问答助手 (Hemiplegia/Stroke Q&A Assistant)",
-    description="由 Qwen-1.8B-Chat LoRA 微调得到的模型 (jinv2/qwen-1_8b-hemiplegia-lora)。与天算AI相关。**医疗建议请咨询专业医生。**",
+    title="偏瘫脑血栓问答助手 (CPU Version - Expect Slow Response)",
+    description=(
+        "由 Qwen-1.8B-Chat LoRA 微调得到的模型 (jinv2/qwen-1_8b-hemiplegia-lora)。与天算AI相关。\n"
+        "**重要：此版本运行在 CPU 上，无量化，响应会非常慢。医疗建议请咨询专业医生。**"
+    ),
     examples=[
         ["偏瘫患者的早期康复锻炼有哪些？"],
         ["什么是脑血栓？"],
         ["中风后如何进行语言恢复训练？"]
     ],
-    allow_flagging="never" # Disable flagging for simplicity
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch() # debug=True can be helpful for local testing but not for Spaces deployment