Update app.py

app.py

@@ -1,35 +1,60 @@
-from fastapi import FastAPI
-import uvicorn
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
+from exllamav2 import (
+    ExLlamaV2,
+    ExLlamaV2Config,
+    ExLlamaV2Cache,
+    ExLlamaV2Tokenizer
+)
+from exllamav2.generator import (
+    ExLlamaV2StreamingGenerator,
+    ExLlamaV2Sampler
+)
+import torch
 
-# Load model and tokenizer
+# Configure model
+model_dir = "TheBloke_Wizard-Vicuna-13B-GPTQ"  # Path to downloaded model
+config = ExLlamaV2Config()
+config.model_dir = model_dir
+config.prepare()
 
-model_id = "DavidAU/Llama-3.2-4X3B-MOE-Hell-California-Uncensored-10B-GGUF"
-#filename = "Llama-3.2-4X3B-MOE-Hell-California-10B-D_AU-Q3_k_s.gguf"
-filename = "Llama-3.2-4X3B-MOE-Hell-California-10B-D_AU-Q2_k.gguf"
+# Load model
+model = ExLlamaV2(config)
+cache = ExLlamaV2Cache(model)
+model.load_autosplit(cache)
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)
-tokenizer.pad_token = tokenizer.eos_token
+# Load tokenizer (HF-compatible)
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
 
-model = AutoModelForCausalLM.from_pretrained(model_id, gguf_file=filename)
-
-def llama2_chat(prompt):
-    inputs = tokenizer(prompt, 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,
-        max_length=30,
-        temperature=0.3
+def generate_response(prompt, max_tokens=200, temperature=0.7):
+    # Initialize generator
+    generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer)
+    generator.set_stop_conditions([tokenizer.eos_token_id])
+    
+    # Configure sampler
+    settings = ExLlamaV2Sampler.Settings()
+    settings.temperature = temperature
+    settings.top_k = 50
+    settings.top_p = 0.8
+    
+    # Encode prompt
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
+    
+    # Generate
+    output = generator.generate_simple(
+        input_ids,
+        settings,
+        max_tokens,
+        seed=42
     )
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
-
-
+    
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
 
 
+##############################################
 
+from fastapi import FastAPI
+import uvicorn
 
 app = FastAPI()
 
@@ -39,7 +64,7 @@ def greet_json():
 
 @app.get("/message")
 async def message(input: str):
-    return llama2_chat(input)
+    return generate_response(input)
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)