Update app.py

app.py

@@ -1,11 +1,13 @@
 from fastapi import FastAPI
 import uvicorn
-from transformers import AutoTokenizer, AutoModel
+from awq import AutoAWQForCausalLM
+from transformers import AutoTokenizer
 
 model_name = "TheBloke/Guanaco-7B-Uncensored-AWQ"
 
-model = AutoModel.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoAWQForCausalLM.from_quantized(model_name_or_path, fuse_layers=True,
+                                          trust_remote_code=False, safetensors=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=False)
 
 app = FastAPI()
 
@@ -15,15 +17,20 @@ def greet_json():
 
 @app.get("/message")
 async def message(input: str):
-    prompt = "### Human: " + input + "\n### Assistant:"
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    prompt=f'''### Human: {input}
+    ### Assistant:
+    
+    '''
+    
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).input_ids.cpu()
     
     output = model.generate(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"],  # Pass attention_mask!
-        max_new_tokens=100,
-        temperature=0.0,  # Disables randomness
-        do_sample=False  # Greedy decoding
+        inputs,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+        top_k=40,
+        max_new_tokens=512
     )
     
     response = tokenizer.decode(output[0], skip_special_tokens=True)