Update app.py

app.py

@@ -8,30 +8,26 @@ import torch
 from transformers import (
     AutoModelForCausalLM,
     BitsAndBytesConfig,
-    AutoTokenizer,
+    GemmaTokenizerFast,
     TextIteratorStreamer,
 )
 
 DESCRIPTION = """\
-# ORLM LLaMA-3-8B
-Hello! I'm ORLM-LLaMA-3-8B, here to automate your optimization modeling tasks! Check our [repo](https://github.com/Cardinal-Operations/ORLM) and [paper](https://arxiv.org/abs/2405.17743)!
+# Gemma 2 9B IT
+Gemma 2 is Google's latest iteration of open LLMs.
+This is a demo of [`google/gemma-2-9b-it`](https://huggingface.co/google/gemma-2-9b-it), fine-tuned for instruction following.
+For more details, please check [our post](https://huggingface.co/blog/gemma2).
+👉 Looking for a larger and more powerful version? Try the 27B version in [HuggingChat](https://huggingface.co/chat/models/google/gemma-2-27b-it).
 """
 
-PROMPT_TEMPLATE = """
-Below is an operations research question. Build a mathematical model and corresponding python code using `coptpy` that appropriately addresses the question.
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
-# Question:
-{Question}
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Response:
-"""
-
-MAX_MAX_NEW_TOKENS = 4096
-DEFAULT_MAX_NEW_TOKENS = 4096
-MAX_INPUT_TOKEN_LENGTH = 2048
-
-model_id = "CardinalOperations/ORLM-LLaMA-3-8B"
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+model_id = "google/gemma-2-9b-it"
+tokenizer = GemmaTokenizerFast.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
@@ -41,7 +37,7 @@ model.config.sliding_window = 4096
 model.eval()
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=90)
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -51,12 +47,20 @@ def generate(
     top_k: int = 50,
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
-    if chat_history != []:
-        return "Sorry, I am an instruction-tuned model and currently do not support chatting. Please try clearing the chat history or refreshing the page to ask a new question."
-
-    prompt = PROMPT_TEMPLATE.replace("{Question}", message).strip()
-    tokenized_example = tokenizer(prompt, return_tensors='pt', max_length=MAX_INPUT_TOKEN_LENGTH, truncation=True)
-    input_ids = tokenized_example.input_ids
+    conversation = []
+    for user, assistant in chat_history:
+        conversation.extend(
+            [
+                {"role": "user", "content": user},
+                {"role": "assistant", "content": assistant},
+            ]
+        )
+    conversation.append({"role": "user", "content": message})
+
+    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     input_ids = input_ids.to(model.device)
 
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
@@ -64,13 +68,12 @@ def generate(
         {"input_ids": input_ids},
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        do_sample=False if temperature == 0.0 else True,
+        do_sample=True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
         num_beams=1,
         repetition_penalty=repetition_penalty,
-        eos_token_id=[tok.eos_token_id],
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
@@ -79,9 +82,6 @@ def generate(
     for text in streamer:
         outputs.append(text)
         yield "".join(outputs)
-    
-    # outputs.append("\n\nI have now attempted to solve the optimization modeling task! Please try executing the code in your environment, making sure it is equipped with `coptpy`.")
-    # yield "".join(outputs)
 
 
 chat_interface = gr.ChatInterface(
@@ -96,44 +96,46 @@ chat_interface = gr.ChatInterface(
         ),
         gr.Slider(
             label="Temperature",
-            minimum=0.0,
+            minimum=0.1,
             maximum=4.0,
             step=0.1,
-            value=0.0,
+            value=0.6,
         ),
         gr.Slider(
             label="Top-p (nucleus sampling)",
             minimum=0.05,
             maximum=1.0,
             step=0.05,
-            value=1.0,
+            value=0.9,
         ),
         gr.Slider(
             label="Top-k",
             minimum=1,
             maximum=1000,
             step=1,
-            value=20,
+            value=50,
         ),
         gr.Slider(
             label="Repetition penalty",
             minimum=1.0,
             maximum=2.0,
             step=0.05,
-            value=1.0,
+            value=1.2,
         ),
     ],
     stop_btn=None,
     examples=[
-        ["A lab has 1000 units of medicinal ingredients to make two pills, a large pill and a small pill. A large pill requires 3 units of medicinal ingredients and 2 units of filler. A small pill requires 2 units of medicinal ingredients and 1 unit of filler. The lab has to make at least 100 large pills. However, since small pills are more popular at least 60% of the total number of pills must be small. How many of each should be made to minimize the total number of filler material needed?"],
-        ["Let's say you're on a mission to create the perfect meal plan for yourself. You're focused on getting the right balance of nutrients without digging too deep into your pockets. You have six different food options to consider: Chicken, Oats, Tofu, Rice, Beef, and Salmon. Each of these foods provides a certain amount of protein, carbohydrates, and calories, and they all come with their own distinct prices.\n\nHere's a detailed breakdown of the nutritional content and cost of each food item:\n\n- Chicken: Delivers 14 grams of protein, a single gram of carbohydrates, and 164 calories at a cost of $6.\n- Oats: Gives you 11 grams of protein, 6 grams of carbohydrates, and 210 calories for just $2.\n- Tofu: Offers 8 grams of protein, 12 grams of carbohydrates, and 98 calories at a cost of $9.\n- Rice: Provides 1 gram of protein, a generous 17 grams of carbohydrates, and 92 calories for $3.\n- Beef: Comes packed with 16 grams of protein, 11 grams of carbohydrates, and 211 calories, priced at $7.\n- Salmon: Brings a hefty 19 grams of protein, 13 grams of carbohydrates, and 211 calories but costs $9.\n\nYou want to make sure your meal plan meets the following nutritional targets: at least 70 grams of protein, 117 grams of carbohydrates, and 1837 calories. Considering these six foods, what is the least amount of money you need to spend to meet these dietary needs?\nRemember, your response should only contain the optimal value of the cost to meet the requirements."], 
-        ["Haus Toys can manufacture and sell toy trucks, toy planes, toy boats, and toy trains. The profit from selling one truck is $5, from one plane is $10, from one boat is $8, and from one train is $7. How many types of toys should Haus Toys manufacture to maximize profit?\n\nThere are 890 units of wood available. Manufacturing one truck requires 12 units of wood, one plane requires 20 units of wood, one boat requires 15 units of wood, and one train requires 10 units of wood.\n\nThere are 500 units of steel available. Manufacturing one plane requires 3 units of steel, one boat requires 5 units of steel, one train requires 4 units of steel, and one truck requires 6 units of steel.\n\nIf Haus Toys manufactures trucks, then they will not manufacture trains.\n\nHowever, if they manufacture boats, they will also manufacture planes.\n\nThe number of toy boats manufactured cannot exceed the number of toy trains manufactured."], 
+        ["Hello there! How are you doing?"],
+        ["Can you explain briefly to me what is the Python programming language?"],
+        ["Explain the plot of Cinderella in a sentence."],
+        ["How many hours does it take a man to eat a Helicopter?"],
+        ["Write a 100-word article on 'Benefits of Open-Source in AI research'"],
     ],
 )
 
 with gr.Blocks(css="style.css", fill_height=True) as demo:
     gr.Markdown(DESCRIPTION)
-    # gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
     chat_interface.render()
 
 if __name__ == "__main__":