Update app.py

app.py

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