fffiloni commited on
Commit
80aa4e5
1 Parent(s): f7b9ef5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import gradio as gr
 
2
  import torch
3
  from transformers import pipeline
4
 
@@ -16,7 +17,11 @@ def infer(user_prompt):
16
  outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
17
  print(outputs)
18
 
19
- return outputs
 
 
 
 
20
 
21
  gr.Interface(
22
  fn = infer,
 
1
  import gradio as gr
2
+ import re
3
  import torch
4
  from transformers import pipeline
5
 
 
17
  outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
18
  print(outputs)
19
 
20
+ pattern = r'\<\|system\|\>(.*?)\<\|assistant\|\>'
21
+ cleaned_text = re.sub(pattern, '', outputs[0]["generated_text"], flags=re.DOTALL)
22
+
23
+
24
+ return cleaned_text
25
 
26
  gr.Interface(
27
  fn = infer,