tiny

app.py

@@ -7,7 +7,7 @@ from models import get_tiny_llama, response_tiny_llama
 from constants import css, js_code, js_light
 
 MERA_table = None
-TINY_LLAMA = get_tiny_llama()
+TINY_LLAMA = None
 
 def giga_gen(content):
     res = generate(content,'auth_token.json')
@@ -74,6 +74,7 @@ if __name__ == "__main__":
     # data_load(args.result_file)    
     # TYPES = ["number", "markdown", "number"]
 
+    TINY_LLAMA =  get_tiny_llama()
     demo = build_demo()
     demo.launch(share=args.share, height=3000, width="110%") # share=args.share
 

models.py

@@ -18,6 +18,6 @@ def response_tiny_llama(
         {"role": "user", "content": content},
     ]
     prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    outputs = pipe(prompt, max_new_tokens=32, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
 
-    return outputs[0]['generated_text']
+    return outputs[0]['generated_text'].split('<|assistant|>')[1]