test_summarization.py

import warnings
warnings.filterwarnings("ignore")
import torch
from transformers import (AutoModelForCausalLM, AutoTokenizer)
from easykv import enable_fixed_kv

# define the model path and the corresponding prompt template
MODEL_CONFIGS = {
    'wizardlm_13b': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/hf_hub/models--WizardLM--WizardLM-13B-V1.2/snapshots/cf5f40382559f19e13874e45b39575171ca46ef8', template="A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n\nUSER: Hello!\nASSISTANT: Hello!</s>\nUSER: {inst}\nASSISTANT:"),
    'llama2_13b_chat': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/hf_hub/models--meta-llama--Llama-2-13b-chat-hf/snapshots/c2f3ec81aac798ae26dcc57799a994dfbf521496/', template="[INST] <<SYS>>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n<</SYS>>\n{inst}[/INST]"),
    'vicuna_13b': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/hf_hub/models--lmsys--vicuna-13b-v1.5/snapshots/3deb0106f72a3a433f0c6ea0cb978bdf14bcd3a6/', template="A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n\nUSER: Hello!\nASSISTANT: Hello!</s>\nUSER: {inst}\nASSISTANT:"),
    'openchat': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/hf_hub/models--openchat--openchat_v3.2_super/snapshots/aab7ce4d48b31a295a0116b61569d8e87a09bb7a/', template="GPT4 User: {inst}<|end_of_turn|>GPT4 Assistant:"),
    'vicuna_7b': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/hf_hub/models--lmsys--vicuna-7b-v1.5/snapshots/de56c35b1763eaae20f4d60efd64af0a9091ebe5/', template="A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n\nUSER: Hello!\nASSISTANT: Hello!</s>\nUSER: {inst}\nASSISTANT:"),
    'wizardlm_7b': dict(path='/cpfs01/user/rensiyu/language_modeling/stanford_alpaca/output_mle_fp16_recycledWiz70k_llama2_7b_512', template="Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{inst}\n\n### Response:"),
    'alpaca_7b': dict(path='/cpfs01/user/rensiyu/language_modeling/stanford_alpaca/output_mle_recycledAlpaca52k_llama2_7b_512_ds', template="Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{inst}\n\n### Response:"),
    'zephyr_7b': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/hf_hub/models--HuggingFaceH4--zephyr-7b-beta/snapshots/dc24cabd13eacd3ae3a5fe574bd645483a335a4a/', template="<|system|>\nYou are a friendly chatbot who always responds in a helpful and detailed manner to the user's questions.</s>\n<|user|>\n{inst}</s>\n<|assistant|>\n"),
    'llama2_7b_chat': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/llama2/model_weights_hf/llama-2-7b-chat-hf/', template="[INST] <<SYS>>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n<</SYS>>\n{inst}[/INST]"),
    'llama2_7b': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/llama2/model_weights_hf/llama-2-7b-hf'),
    'llama2_13b': dict(path='/cpfs01/shared/public/public_hdd/llmeval/model_weights/hf_hub/models--meta-llama--Llama-2-13b-hf/snapshots/dc1d3b3bfdb69df26f8fc966c16353274b138c55/'),
}

model_name = 'llama2_7b_chat'
path = MODEL_CONFIGS[model_name]['path']
template = MODEL_CONFIGS[model_name]['template']
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16, device_map='auto').eval()
tokenizer = AutoTokenizer.from_pretrained(path)

# Setup KV caching mode
stride = 24
enable_fixed_kv(model, tokenizer, mode='encoding', stride=stride)

# Test input
article = "###\nArticle: Media playback is not supported on this device The tournament kicks off with France playing Romania on Friday, 10 June and the final will take place at the Stade de France in Paris on 10 July. Spain are hoping to retain the title they won in 2008 and 2012. England will be appearing at their ninth finals, while Wales and Northern Ireland have qualified for the first time. Good question - because, for the first time, there are 24 teams competing in the finals. That is an increase from the 16 that had taken part in every edition since it was hosted in England in 1996. With six groups of four teams, it means the top two will qualify for the last 16, plus the four best third-placed finishers. In other words, only eight teams will fail to qualify from the group stage. One point could be enough to put your team into the last 16 - and from then the tournament goes to a knockout format. As Greece proved by pulling off a sensational triumph at Euro 2004, trying to predict a winner is a difficult game. World champions Germany are understandably one of the favourites, while holders Spain are also short odds with the bookmakers. Germany finished top of their qualification group but booking their place in France was not without its hiccups - with defeats by Poland and the Republic of Ireland. They have been beaten by France and England since qualifying but remain one to watch. Spain dominated international football between 2008 and 2012 but failed to qualify from their group at the 2014 World Cup. Could this be a last hurrah for Vicente del Bosque's ageing but brilliant side? France won the World Cup as hosts in 1998 - can they repeat the famous success of the team led by Didier Deschamps, Zinedine Zidane and Laurent Blanc? A run into the latter stages by Les Bleus is likely to help restore national morale in a country still recovering from last year's deadly Paris attacks. France are the bookmakers' favourites to win the whole thing, but their form is tricky to gauge because they did not have to qualify for this tournament. Fast on the counter and unified after years of internal division, watch out for talented midfielders Paul Pogba and Blaise Matuidi. But they cannot call upon prolific Real Madrid striker Karim Benzema. The French Football Federation (FFF) said he would not be picked after being investigated for his part in an alleged plot to blackmail Les Bleus team-mate Mathieu Valbuena, who did not make the squad. Well, there are three of them in France - only Scotland missed out - plus the Republic of Ireland. It is going to get crowded too, because England and Wales are both in Group B, along with Russia and debutants Slovakia. Don't expect much work to get done when England and Wales play each other - that match in Lens will be shown live on BBC One, kick-off 14:00 BST, on Thursday, 16 June. Roy Hodgson's England have a good young team but doubts persist about their defence. They breezed through their qualifying group with a perfect record of 10 wins, while confidence was further boosted with March's impressive win away at world champions Germany. Wales have got Real Madrid star Gareth Bale and plenty of heart. Chris Coleman's side reached their first major tournament in 57 years after losing just once in qualification. Northern Ireland might have numerous players from England's lower leagues but they finished top of their qualifying group, losing just one of their 10 matches. Striker Kyle Lafferty could not get a game for club side Norwich during qualifying (he has since been loaned to Birmingham) but he scored seven crucial goals in nine games for his country. With the 24-team format offering hope to traditionally less-successful nations, five countries have taken advantage to secure their European Championship finals debut. Northern Ireland and Wales, of course, are two of them. Slovakia, who have been drawn alongside England and Wales in Group B, have qualified for the first time as an independent state. A 1-0 win against Spain showed the 2010 World Cup qualifiers can mix it with Europe's elite. Beware England and Wales. Albania have never been near a major finals, but edged out Denmark - the 1992 champions - in their qualifying group thanks to a shock win in Portugal. Oh, and because they were awarded a 3-0 win in Serbia by the Court of Arbitration for Sport after a riot. With a population of roughly 330,000 (comparable to Coventry) and only 21,508 registered players, Iceland are unsurprisingly the smallest nation to ever qualify for the finals. How did they do that? By beating the Netherlands - 1988 winners and three-time semi-finalists - home and away during qualifying. The Dutch, for the first time since 1984, will not be there. Real Madrid superstar Cristiano Ronaldo, who was voted the world's second best player behind Argentina's Lionel Messi in the 2015 Ballon d'Or vote, heads a star-studded cast. Ronaldo will be leading Portugal's challenge, while Wales hope his club-mate Gareth Bale - the world's most expensive player - can transform his impressive La Liga form onto the international stage. World Cup winners Germany boast a host of stellar names who have impressed in previous tournaments, most notably Bayern Munich goalkeeper Manuel Neuer and forward Thomas Muller plus Real Madrid midfielder Toni Kroos. France's challenge is set to be driven by energetic Juventus midfielder Paul Pogba, who is still reportedly courting the attention of several English Premier League sides. Belgium are not short of star quality either. Premier League pair Eden Hazard and Kevin de Bruyne will spearhead their quest for a first major tournament win, but the Red Devils are without injured captain Vincent Kompany. And Sweden superstar Zlatan Ibrahimovic, 34, will be hoping to illuminate an international tournament for what could be the final time. The Netherlands' failure to qualify leaves Bayern Munich winger Arjen Robben watching from home, while his club-mate Franck Ribery was not named in the France squad, despite suggestions he was considering ending his international retirement. The tournament's 51 games will be staged at 10 locations across France, including new stadiums in Bordeaux, Lille, Lyon and Nice. The opening match - between France and Romania on 10 June - and the 10 July final will be played at the Stade de France in Paris. Building the new venues and renovating historic grounds such as Marseille's Stade Velodrome has cost 1.6bn euros (\u00a31.2bn) - modernisation which was necessary, organisers say, because France did not fully capitalise on hosting the 1998 World Cup. Vibrating stands, floating roofs and adjacent slag heaps - read BBC Sport's venue-by-venue guide In March this year, the International Football Association Board (IFAB) ratified a host of revisions to the laws of the game in an attempt to remove inconsistencies and meet the needs of the modern game. The changes came into effect on 1 June, so will apply for Euro 2016. More than 90 revisions were made, but these are some of the key changes: Kick-off: Previously, the ball had to go forward from kick-off but the rule has been changed to allow it to go in any direction. Pre-match red cards: Players can now be sent off before a match gets under way, although they can be replaced by another player in the match-day squad. Leaving the pitch after treatment: If a player is injured in a challenge resulting in a yellow or red card, they no longer have to leave the field and can have a quick assessment or medical treatment. This change is designed to prevent situations where a team would be temporarily down to 10 players. The end of 'triple punishment': A professional foul inside the area will now normally result in a yellow card for the offender, and not a red. This is to end the so-called triple punishment of penalty, dismissal and suspension, which was seen by some as excessive. There are exceptions for when the offender will receive a red which include holding, pushing or pulling and violent conduct. Goalline technology, already established in the Premier League and elsewhere, will be used at Euro 2016. Yes. But you had better be quick. Most of the group games - including all of England's matches - are sold out. However, as of 9 June there are tickets available to watch Wales and Northern Ireland, plus other games involving some of the smaller nations. There are also tickets available for the opening game between France and Romania - at 395 euros each. Find the latest ticket details on Uefa's official Euro 2016 website Each of the 10 host cities will have official Uefa fan zones. The fan zones will have a giant screen for showing all of the matches throughout the tournament. The biggest will be in Paris, where up to 90,000 supporters can gather on the Champ de Mars - in the shadow of the Eiffel Tower. You won't miss a kick. The BBC and ITV will bring audiences closer to the heart of the action than ever before with extensive coverage of Euro 2016. The BBC will showcase 26 live matches on TV and streamed online within its 140 hours of programming, including highlights on BBC One, BBC Two and the Red Button plus 24 days of coverage on Radio 5 live. The BBC's Euro 2016 digital service will feature live broadcasts of all BBC games, commentary and live text coverage of every game featuring in-game highlight clips, on-demand highlights of every goal, breaking news, exclusive features and expert analysis. ITV will be home to 26 live games across ITV1 and ITV4, which will be simulcast live on the ITV Hub. READ MORE: Ferdinand and Henry join BBC coverage for the Euros The French government has extended a state of emergency imposed after the Paris attacks in November to cover Euro 2016. The coordinated gun and bomb attacks in Paris on 13 November killed 130 people and were claimed by so-called Islamic State. The Stade de France, which was targeted by suicide bombers, will host the opening match of Euro 2016 and the final. The current state of emergency gives police extra powers to conduct searches and put people under house arrest. More than 90,000 police, soldiers and private security agents are being deployed as well. US warns of Euro 2016 terror threat\n\nSummarize the above article in 1 sentence.\n"
prompt = f"Write a SHORT summary of the following text delimited by triple backticks. Return your response which covers the key points of the text.\n```{article}```"
input_prompt = template.format(inst=prompt)

# Define eviction policy
kv_policy = 'roco'

for budget in [1.0, 0.5]:
    # Define sampling parameters
    gen_kwargs = dict(
        temperature=0.3,
        top_p=1.0,
        max_new_tokens=256,
        budget=budget,
        kv_policy=kv_policy,
        keep_attention=True
    )
    input_ids = tokenizer([input_prompt], return_tensors='pt').input_ids.to(model.device)
    output = model.easykv_generate(input_ids=input_ids, generation_config=gen_kwargs)
    print(f"{'='*20} EasyKV-{kv_policy} {'='*20}\n{output}")