70b v1

app.py

@@ -5,11 +5,10 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-
 DESCRIPTION = """\
-# Qwen2 0.5B Instruct Text Completion
+# Llama-3.1-70B-Instruct Text Completion
 
-This is a demo of [`Qwen/Qwen2-0.5B-Instruct`](https://huggingface.co/Qwen/Qwen2-0.5B-Instruct), fine-tuned for instruction following.
+This is a demo of [`meta-llama/Llama-3.1-70B-Instruct`](https://huggingface.co/meta-llama/Llama-3.1-70B-Instruct), a large language model fine-tuned for instruction following.
 
 Enter your text in the box below and click "Complete" to have the AI generate a completion for your input. The generated text will be appended to your input. You can stop the generation at any time by clicking the "Stop" button.
 """
@@ -19,26 +18,26 @@ DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-model_id = "Qwen/Qwen2-0.5B-Instruct"
+model_id = "meta-llama/Llama-3.1-70B-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.float16,
+    load_in_8bit=True,
 )
 model.eval()
-
 @spaces.GPU(duration=90)
 def generate(
     message: str,
-    max_new_tokens: int = 420,
-    temperature: float = 0,
+    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]:
-    input_ids = tokenizer.encode(message, return_tensors="pt")
+    prompt = f"[INST] {message} [/INST]"
+    input_ids = tokenizer.encode(prompt, 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 as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
@@ -90,7 +89,7 @@ with gr.Blocks(css="style.css", fill_height=True) as demo:
                 minimum=0.1,
                 maximum=1.0,
                 step=0.1,
-                value=0,
+                value=0.6,
             )
             top_p = gr.Slider(
                 label="Top-p (nucleus sampling)",