Update proton_code_llama_fine_tuning_using_qlora_inferencing.py

proton_code_llama_fine_tuning_using_qlora_inferencing.py

@@ -105,19 +105,46 @@ def prepare_prompt(query):
     prompt_prep = ''.join([priming_text,getExamples(query),'\n',decorator,'\n','Similar to the above examples, generate only one Proton UI for the following description:',query])
     return prompt_prep
 
-def generate(query):
-    prompt_prep = prepare_prompt(query)
-    encodings = model_tokenizer(prompt_prep, return_tensors = 'pt').to(device)
-    
-    with torch.inference_mode():
-        outputs = model.generate(input_ids=encodings.input_ids,
-                            attention_mask = encodings.attention_mask,
-                            generation_config = generation_config,
+def generate_text(query):
+  prompt_prep = prepare_prompt(query)
+  encodings = model_tokenizer(prompt_prep, return_tensors = 'pt').to(device)
+
+  # with torch.inference_mode():
+  outputs = model.generate(input_ids=encodings.input_ids,
+                          attention_mask = encodings.attention_mask,
+                          generation_config = generation_config,
                           max_new_tokens=1000,do_sample=True)
+
+  outputs = model_tokenizer.decode(outputs[0],skip_special_tokens=True,best_of=1)
+  outputs = outputs[outputs.find(query): ]
+  return re.match(r'(.*?\nDescription:.*?\nDescription:)', outputs, re.DOTALL).group(1)[:-12]
+
+def alternatingly_agree(message, history):
+  output = generate_text(message)
+  return output
+# gr.ChatInterface(alternatingly_agree).launch(inline=False,enable_queue=True)
+
+gr.ChatInterface(
+    alternatingly_agree,
+    retry_btn=None,
+    undo_btn="Delete Previous",
+    clear_btn="Clear",
+).launch(inline=False,
+    enable_queue=True,share=True)
+
+# def generate(query):
+#     prompt_prep = prepare_prompt(query)
+#     encodings = model_tokenizer(prompt_prep, return_tensors = 'pt').to(device)
+    
+#     with torch.inference_mode():
+#         outputs = model.generate(input_ids=encodings.input_ids,
+#                             attention_mask = encodings.attention_mask,
+#                             generation_config = generation_config,
+#                           max_new_tokens=1000,do_sample=True)
     
-    outputs = model_tokenizer.decode(outputs[0],skip_special_tokens=True,best_of=1)
+#     outputs = model_tokenizer.decode(outputs[0],skip_special_tokens=True,best_of=1)
     
-    return outputs[outputs.find(query): ]#outputs.find('Explanation')
+#     return outputs[outputs.find(query): ]#outputs.find('Explanation')
 
 # output = generate('Create a choice type of Proton UI for choosing vegetables')
 # print(output)