Update app.py

app.py

@@ -10,7 +10,11 @@ if device.type == "cpu":
 # Load the tokenizer and model directly
 model_name = "ruslanmv/ai-medical-model-32bit"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+model = model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map="auto", 
+    load_in_8bit=True
+).to(device)
 
 
 
@@ -21,11 +25,11 @@ def ask_medical_question(question):
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     outputs = model.generate(
         **inputs,
-        max_new_tokens=256,
+        max_new_tokens=100,
         temperature=0.7,
         do_sample=True,
-        top_p=0.95,
-        top_k=50,
+        top_p=0.9,
+        top_k=30,
     )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response