Update app.py

app.py

@@ -1,22 +1,11 @@
 from fastapi import FastAPI
 import uvicorn
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_name = "TheBloke/Wizard-Vicuna-13B-Uncensored-HF"
+model_name = "Tap-M/Luna-AI-Llama2-Uncensored"
 
-# Configure 4-bit quantization
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,               # Enable 4-bit quantization
-    bnb_4bit_quant_type="nf4",       # Use 4-bit NormalFloat (optimal)
-    bnb_4bit_compute_dtype="float16", # Faster computation with float16
-    bnb_4bit_use_double_quant=True,   # Extra compression
-    llm_int8_enable_fp32_cpu_offload=True
-)
-
-# Load model with quantization
 model = AutoModelForCausalLM.from_pretrained(
     model_name,  # Example model
-    quantization_config=bnb_config,
     device_map="auto",            # Auto-distribute across GPU/CPU
     trust_remote_code=True        # Required for some models
 )
@@ -33,7 +22,9 @@ def greet_json():
 
 @app.get("/message")
 async def message(input: str):
-    inputs = tokenizer(input, return_tensors="pt", padding=True, truncation=True)
+    prompt = "USER:" + input + "\nASSISTANT:"
+    
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
     
     output = model.generate(
         input_ids=inputs["input_ids"],