Update app.py

app.py

@@ -29,7 +29,7 @@ if torch.cuda.is_available():
     model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.padding_side = 'right'
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    # pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
     # tokenizer.use_default_system_prompt = False
     
 @spaces.GPU(duration=60)
@@ -50,38 +50,39 @@ def generate(
         conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
     conversation.append({"role": "user", "content": message})
 
-    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,
+#    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.