Update app.py

app.py

@@ -1,18 +1,65 @@
 from fastapi import FastAPI
 import uvicorn
-from transformers import AutoModel, AutoTokenizer
 
-model_name = "Tap-M/Luna-AI-Llama2-Uncensored"
+#model_name = "Tap-M/Luna-AI-Llama2-Uncensored"
 
-model = AutoModel.from_pretrained(
-    model_name,  # Example model
-    offload_folder="./offload",  # Temporary directory
-    trust_remote_code=True        # Required for some models
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer
+import torch
+
+# Configuration for 4-bit quantization
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",        # Optimized 4-bit precision
+    bnb_4bit_compute_dtype=torch.float16,  # Faster computations
+    bnb_4bit_use_double_quant=True    # Extra memory savings
+)
+
+# Load model and tokenizer
+model_name = "meta-llama/Llama-2-7b-chat-hf"  # or "13b-chat-hf"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=bnb_config,
+    device_map="auto",                # Auto-distribute across GPU/CPU
+    torch_dtype=torch.float16,
+    trust_remote_code=True            # Required for Llama 2
 )
 
-# load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token
+# Set chat template (critical for chat models)
+tokenizer.chat_template = "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + content + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + content + ' ' + eos_token }}{% endif %}{% endfor %}"
+
+def llama2_chat(prompt, system_prompt="You are a helpful assistant."):
+    # Format as Llama 2 chat
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": prompt}
+    ]
+    
+    # Tokenize with chat template
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        return_tensors="pt"
+    ).to(model.device)
+    
+    # Stream output tokens
+    streamer = TextStreamer(tokenizer, skip_prompt=True)
+    
+    # Generate response
+    outputs = model.generate(
+        inputs,
+        max_new_tokens=1000,
+        temperature=0.7,
+        streamer=streamer
+    )
+    
+    # Decode full output
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+
+
+
+
 
 app = FastAPI()
 
@@ -22,19 +69,7 @@ def greet_json():
 
 @app.get("/message")
 async def message(input: str):
-    prompt = "USER:" + input + "\nASSISTANT:"
-    
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-    
-    output = model.generate(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"], 
-        max_new_tokens=100,
-    )
-    
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    
-    return response
+    return llama2_chat(input)
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)