Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import torch | |
TOKENIZER_REPO = "MediaTek-Research/Breeze-7B-Instruct-v1_0" | |
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_REPO,local_files_only=False,use_fast=True) | |
model = AutoModelForCausalLM.from_pretrained( | |
TOKENIZER_REPO, | |
device_map="auto", | |
local_files_only=False, | |
torch_dtype=torch.bfloat16 | |
) | |
def generate(text): | |
chat_data = [] | |
text = text.strip() | |
print("text===="+text) | |
if text: | |
chat_data.append({"role": "system", "content": text}) | |
print(chat_data) | |
achat=tokenizer.apply_chat_template(chat_data,return_tensors="pt") | |
print(achat) | |
outputs = model.generate(achat, | |
max_new_tokens=128, | |
top_p=0.01, | |
top_k=85, | |
repetition_penalty=1.1, | |
temperature=0.01) | |
theResult=tokenizer.decode(outputs[0]) | |
print(theResult) | |
splitOutput=theResult.splitlines() | |
theReturn="" | |
for i in range(0,len(splitOutput)): | |
print("i={},out={}".format(i, splitOutput[i])) | |
if(i>0 and splitOutput[i].strip()): | |
theReturn+=splitOutput[i].strip() | |
print("result={}".format(theReturn)) | |
return tokenizer.decode(outputs[0]) | |
gradio_app = gr.Interface( | |
generate, | |
inputs=gr.Text(), | |
outputs=gr.Text(), | |
title="test", | |
) | |
if __name__ == "__main__": | |
gradio_app.launch() | |