Update app.py

app.py

@@ -19,26 +19,6 @@ conversations = {}
 Device_Type = "cuda"
 
 
-def load_quantized_model(model_id, model_basename):
-    # The code supports all huggingface models that ends with GPTQ and have some variation
-    # of .no-act.order or .safetensors in their HF repo.
-    print("Using AutoGPTQForCausalLM for quantized models")
-
-    if ".safetensors" in model_basename:
-        # Remove the ".safetensors" ending if present
-        model_basename = model_basename.replace(".safetensors", "")
-
-    quantized_tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
-    print("Tokenizer loaded")
-
-    quantized_model = AutoGPTQForCausalLM.from_quantized(model_id, model_basename=model_basename, use_safetensors=True, trust_remote_code=True, device_map="auto", use_triton=False, quantize_config=None,)
-
-    return quantized_model, quantized_tokenizer
-
-
-# Making the code device-agnostic
-#model, tokenizer = load_quantized_model(model_name_or_path, "model.safetensors")
-
 def load_model_norm():
     if torch.cuda.is_available():
         print("CUDA is available. GPU will be used.")
@@ -50,7 +30,7 @@ def load_model_norm():
     # For example: revision="main"
     model = AutoModelForCausalLM.from_pretrained(model_name_or_path,device_map="auto", trust_remote_code=True,revision="gptq-4bit-128g-actorder_True")
     # Switch to CPU inference
-    model.to("cuda")
+    #model.to("cuda")
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
 
     return model, tokenizer 
@@ -91,9 +71,18 @@ def generate_response(prompt: str) -> str:
 
     prompt_template = f'{PERSONA_DESC}\n\nASSISTANT: {prompt}\n'
 
-    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
-    output = model.generate(inputs=input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
-    generated_text = tokenizer.decode(output[0])
+    pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=512,
+    do_sample=True,
+    temperature=0.7,
+    top_p=0.95,
+    top_k=40,
+    repetition_penalty=1.1
+)
+    generated_text = (pipe(prompt_template)[0]['generated_text'])
     
 
     return generated_text