Update app.py

app.py

@@ -28,9 +28,10 @@ if torch.cuda.is_available():
     model_id = "patched-codes/patched-mix-4x7B"
     model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    tokenizer.use_default_system_prompt = False
-
-
+    tokenizer.padding_side = 'right'
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    # tokenizer.use_default_system_prompt = False
+    
 @spaces.GPU(duration=60)
 def generate(
     message: str,
@@ -49,33 +50,38 @@ def generate(
         conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-    input_ids = input_ids.to(model.device)
-
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        #top_k=top_k,
-        temperature=temperature,
-        eos_token_id=tokenizer.eos_token_id, 
-        pad_token_id=tokenizer.pad_token_id,
+    prompt = pipe.tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    outputs = pipe(prompt, max_new_tokens=max_new_tokens, do_sample=True, temperature=temperature, top_p=top_p, 
+                   eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)
+
+    return outputs[0]['generated_text'][len(prompt):].strip()
+#    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
+#    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+#        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+#        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+#    input_ids = input_ids.to(model.device)
+
+#    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+#    generate_kwargs = dict(
+#        {"input_ids": input_ids},
+#        streamer=streamer,
+#        max_new_tokens=max_new_tokens,
+#        do_sample=True,
+#        top_p=top_p,
+#        #top_k=top_k,
+#        temperature=temperature,
+#        eos_token_id=tokenizer.eos_token_id, 
+#        pad_token_id=tokenizer.pad_token_id,
         #num_beams=1,
         #repetition_penalty=1.2,
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+#    )
+#    t = Thread(target=model.generate, kwargs=generate_kwargs)
+#    t.start()
+
+#    outputs = []
+#    for text in streamer:
+#        outputs.append(text)
+#        yield "".join(outputs)
 
 example1='''You are a senior software engineer who is best in the world at fixing vulnerabilities.
 Users will give you vulnerable code and you will generate a fix based on the provided INSTRUCTION.