Update app.py

app.py

@@ -1,54 +1,29 @@
 from fastapi import FastAPI
 import uvicorn
 
-model_name = "TheBloke/Luna-AI-Llama2-Uncensored-GPTQ"
+model_name = "Llama-3.2-4X3B-MOE-Hell-California-Uncensored-10B-GGUF"
 
-from transformers import AutoModel, AutoTokenizer, TextStreamer, BitsAndBytesConfig
+from transformers import AutoModel, AutoTokenizer, TextStreamer
 import torch
 
-bnb_config = BitsAndBytesConfig(
-    load_in_8bit=True,
-    llm_int8_enable_fp32_cpu_offload=True,
-)
-
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModel.from_pretrained(
     model_name,
-    quantization_config=bnb_config,
-    device_map="cpu",
-    trust_remote_code=True            # Required for Llama 2
+    device_map="auto",
+    trust_remote_code=True
 )
 
-# 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
+def llama2_chat(prompt):
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)    
+    output = model.generate(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],  # Pass attention_mask!
+        max_new_tokens=100,
+        temperature=0.3
     )
-    
-    # Decode full output
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response