Update app.py

app.py

@@ -1,172 +1,138 @@
-import spaces
 import os
-import json
-import subprocess
-from llama_cpp import Llama
-from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
-from llama_cpp_agent.providers import LlamaCppPythonProvider
-from llama_cpp_agent.chat_history import BasicChatHistory
-from llama_cpp_agent.chat_history.messages import Roles
-import gradio as gr
-from huggingface_hub import hf_hub_download
+from threading import Thread
+from typing import Iterator
 
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, GemmaTokenizerFast, TextIteratorStreamer
+
+DESCRIPTION = """\
+# Gemma 2 2B IT
+Gemma 2 is Google's latest iteration of open LLMs.
+This is a demo of [`google/gemma-2-2b-it`](https://huggingface.co/google/gemma-2-2b-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) and the 9B version in [this Space](https://huggingface.co/spaces/huggingface-projects/gemma-2-9b-it).
+"""
 
-huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-
-hf_hub_download(
- repo_id="bartowski/gemma-2-9b-it-GGUF",
- filename="gemma-2-9b-it-Q5_K_M.gguf",
- local_dir="./models"
-)
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
-hf_hub_download(
- repo_id="bartowski/gemma-2-27b-it-GGUF",
- filename="gemma-2-27b-it-Q5_K_M.gguf",
- local_dir="./models"
-)
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-hf_hub_download(
- repo_id="google/gemma-2-2b-it-GGUF",
- filename="2b_it_v2.gguf",
- local_dir="./models",
- token=huggingface_token
+model_id = "google/gemma-2-2b-it"
+tokenizer = GemmaTokenizerFast.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+ model_id,
+ device_map="auto",
+ torch_dtype=torch.bfloat16,
 )
-
-
-
-llm = None
-llm_model = None
-
-@spaces.GPU(duration=120)
-def respond(
- message,
- history: list[tuple[str, str]],
- model,
- system_message,
- max_tokens,
- temperature,
- top_p,
- top_k,
- repeat_penalty,
-):
- chat_template = MessagesFormatterType.GEMMA_2
-
- global llm
- global llm_model
- 
- if llm is None or llm_model != model:
- llm = Llama(
- model_path=f"models/{model}",
- flash_attn=True,
- n_gpu_layers=81,
- n_batch=1024,
- n_ctx=8192,
+model.config.sliding_window = 4096
+model.eval()
+
+
+@spaces.GPU(duration=90)
+def generate(
+ message: str,
+ chat_history: list[tuple[str, str]],
+ max_new_tokens: int = 1024,
+ temperature: float = 0.6,
+ top_p: float = 0.9,
+ top_k: int = 50,
+ repetition_penalty: float = 1.2,
+) -> Iterator[str]:
+ conversation = []
+ for user, assistant in chat_history:
+ conversation.extend(
+ [
+ {"role": "user", "content": user},
+ {"role": "assistant", "content": assistant},
+ ]
  )
- llm_model = model
-
- provider = LlamaCppPythonProvider(llm)
-
- agent = LlamaCppAgent(
- provider,
- system_prompt=f"{system_message}",
- predefined_messages_formatter_type=chat_template,
- debug_output=True
+ 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)
+ generate_kwargs = dict(
+ {"input_ids": input_ids},
+ streamer=streamer,
+ max_new_tokens=max_new_tokens,
+ do_sample=True,
+ top_p=top_p,
+ top_k=top_k,
+ temperature=temperature,
+ num_beams=1,
+ repetition_penalty=repetition_penalty,
  )
- 
- settings = provider.get_provider_default_settings()
- settings.temperature = temperature
- settings.top_k = top_k
- settings.top_p = top_p
- settings.max_tokens = max_tokens
- settings.repeat_penalty = repeat_penalty
- settings.stream = True
-
- messages = BasicChatHistory()
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
+ t.start()
 
- for msn in history:
- user = {
- 'role': Roles.user,
- 'content': msn[0]
- }
- assistant = {
- 'role': Roles.assistant,
- 'content': msn[1]
- }
- messages.add_message(user)
- messages.add_message(assistant)
- 
- stream = agent.get_chat_response(
- message,
- llm_sampling_settings=settings,
- chat_history=messages,
- returns_streaming_generator=True,
- print_output=False
- )
- 
- outputs = ""
- for output in stream:
- outputs += output
- yield outputs
+ outputs = []
+ for text in streamer:
+ outputs.append(text)
+ yield "".join(outputs)
 
-description = """<p align="center">Defaults to 2B (you can switch to 9B or 27B from additional inputs)</p>
-<p><center>
-<a href="https://huggingface.co/google/gemma-2-27b-it" target="_blank">[27B it Model]</a>
-<a href="https://huggingface.co/google/gemma-2-9b-it" target="_blank">[9B it Model]</a>
-<a href="https://huggingface.co/google/gemma-2-2b-it" target="_blank">[2B it Model]</a>
-<a href="https://huggingface.co/bartowski/gemma-2-27b-it-GGUF" target="_blank">[27B it Model GGUF]</a>
-<a href="https://huggingface.co/bartowski/gemma-2-9b-it-GGUF" target="_blank">[9B it Model GGUF]</a>
-<a href="https://huggingface.co/google/gemma-2-2b-it-GGUF" target="_blank">[2B it Model GGUF]</a>
-</center></p>
-"""
 
-demo = gr.ChatInterface(
- respond,
+chat_interface = gr.ChatInterface(
+ fn=generate,
  additional_inputs=[
- gr.Dropdown([
- 'gemma-2-9b-it-Q5_K_M.gguf',
- 'gemma-2-27b-it-Q5_K_M.gguf',
- '2b_it_v2.gguf'
- ],
- value="2b_it_v2.gguf",
- label="Model"
+ gr.Slider(
+ label="Max new tokens",
+ minimum=1,
+ maximum=MAX_MAX_NEW_TOKENS,
+ step=1,
+ value=DEFAULT_MAX_NEW_TOKENS,
  ),
- gr.Textbox(value="You are a helpful assistant.", label="System message"),
- gr.Slider(minimum=1, maximum=4096, value=2048, step=1, label="Max tokens"),
- gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
  gr.Slider(
+ label="Temperature",
  minimum=0.1,
+ maximum=4.0,
+ step=0.1,
+ value=0.6,
+ ),
+ gr.Slider(
+ label="Top-p (nucleus sampling)",
+ minimum=0.05,
  maximum=1.0,
- value=0.95,
  step=0.05,
- label="Top-p",
+ value=0.9,
  ),
  gr.Slider(
- minimum=0,
- maximum=100,
- value=40,
- step=1,
  label="Top-k",
+ minimum=1,
+ maximum=1000,
+ step=1,
+ value=50,
  ),
  gr.Slider(
- minimum=0.0,
- maximum=2.0,
- value=1.1,
- step=0.1,
  label="Repetition penalty",
+ minimum=1.0,
+ maximum=2.0,
+ step=0.05,
+ value=1.2,
  ),
  ],
- retry_btn="Retry",
- undo_btn="Undo",
- clear_btn="Clear",
- submit_btn="Send",
- title="Chat with Gemma 2 using llama.cpp", 
- description=description,
- chatbot=gr.Chatbot(
- scale=1, 
- likeable=False,
- show_copy_button=True
- )
+ stop_btn=None,
+ examples=[
+ ["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'"],
+ ],
+ cache_examples=False,
 )
 
+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")
+ chat_interface.render()
+
 if __name__ == "__main__":
- demo.launch()
+ demo.queue(max_size=20).launch()