Spaces:
Runtime error
Runtime error
| import os | |
| import sys | |
| import uvicorn | |
| import torch | |
| from fastapi import Body, FastAPI | |
| from transformers import T5Tokenizer, MT5ForConditionalGeneration | |
| import pytorch_lightning as pl | |
| sys.path.append(os.path.abspath(os.path.join( | |
| os.path.dirname(__file__), os.path.pardir))) | |
| os.environ["CUDA_VISIBLE_DEVICES"] = '5' | |
| os.environ["MASTER_ADDR"] = '127.0.0.1' | |
| os.environ["MASTER_PORT"] = '6000' | |
| device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| print('device') | |
| pretrain_model_path = '/cognitive_comp/ganruyi/hf_models/google/mt5-large' | |
| # pretrain_model_path = 'google/mt5-small' | |
| model_path = '/cognitive_comp/ganruyi/fengshen/mt5_large_summary/ckpt/epoch-0-last.ckpt' | |
| tokenizer = T5Tokenizer.from_pretrained(pretrain_model_path) | |
| print('load tokenizer') | |
| class MT5FinetuneSummary(pl.LightningModule): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = MT5ForConditionalGeneration.from_pretrained(pretrain_model_path) | |
| model = MT5FinetuneSummary.load_from_checkpoint(model_path) | |
| print('load checkpoint') | |
| model.to(device) | |
| model.eval() | |
| app = FastAPI() | |
| print('server start') | |
| # def flask_gen(text: str, level: float = 0.9, n_sample: int = 5, length: int = 32, is_beam_search=False): | |
| async def flask_gen(text: str = Body('', title='原文', embed=True), | |
| n_sample: int = 5, length: int = 32, is_beam_search=False): | |
| if len(text) > 128: | |
| text = text[:128] | |
| text = 'summary:'+text | |
| print(text) | |
| # inputs = tokenizer(text, return_tensors='pt') | |
| inputs = tokenizer.encode_plus( | |
| text, max_length=128, padding='max_length', truncation=True, return_tensors='pt') | |
| # print(inputs) | |
| if is_beam_search: | |
| generated_ids = model.model.generate( | |
| input_ids=inputs['input_ids'].to(device), | |
| attention_mask=inputs['attention_mask'].to(device), | |
| max_length=length, | |
| num_beams=n_sample, | |
| repetition_penalty=2.5, | |
| length_penalty=1.0, | |
| early_stopping=True, | |
| num_return_sequences=n_sample | |
| ) | |
| else: | |
| generated_ids = model.model.generate( | |
| input_ids=inputs['input_ids'].to(device), | |
| attention_mask=inputs['attention_mask'].to(device), | |
| max_length=length, | |
| do_sample=True, | |
| temperature=1.0, | |
| top_p=1.0, | |
| repetition_penalty=2.5, | |
| # early_stopping=True, | |
| num_return_sequences=n_sample | |
| ) | |
| result = [] | |
| # print(tokenizer.all_special_tokens) | |
| for sample in generated_ids: | |
| preds = [tokenizer.decode(sample, skip_special_tokens=True, | |
| clean_up_tokenization_spaces=True)] | |
| preds = ''.join(preds) | |
| # print(preds) | |
| result.append(preds) | |
| return result | |
| if __name__ == '__main__': | |
| uvicorn.run(app, host="0.0.0.0", port=6607, log_level="debug") | |
| # # article = "日前，方舟子发文直指林志颖旗下爱碧丽推销假保健品，引起哗然。调查发现， | |
| # 爱碧丽没有自己的生产加工厂。其胶原蛋白饮品无核心研发，全部代工生产。号称有“逆生长”功效的爱碧丽“梦幻奇迹限量组”售价>高达1080元，实际成本仅为每瓶4元！" | |
| # article = '''在北京冬奥会自由式滑雪女子坡面障碍技巧决赛中，中国选手谷爱凌夺得银牌。祝贺谷爱凌！ | |
| # 今天上午，自由式滑雪女子坡面障碍技巧决赛举行。决赛分三轮进行，取选手最佳成绩排名决出奖牌。 | |
| # 第一跳，中国选手谷爱凌获得69.90分。在12位选手中排名第三。完成动作后，谷爱凌又扮了个鬼脸，甚是可爱。 | |
| # 第二轮中，谷爱凌在道具区第三个障碍处失误，落地时摔倒。获得16.98分。网友：摔倒了也没关系，继续加油！ | |
| # 在第二跳失误摔倒的情况下，谷爱凌顶住压力，第三跳稳稳发挥，流畅落地！获得86.23分！此轮比赛，共12位选手参赛，谷爱凌第10位出场。网友：看比赛时我比谷爱凌紧张，加油！''' | |
| # flask_gen(article, length=30) | |