Update app.py

app.py

@@ -5,7 +5,7 @@ from typing import Iterator
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, LlamaForCausalLM
 from peft import AutoPeftModelForCausalLM
 
 DESCRIPTION = """\
@@ -22,9 +22,9 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 HF_TOKEN = os.getenv("HF_TOKEN")
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-model_name = "ehristoforu/BigFalcon3-from10B"
+model_name = "estrogen/c4ai-command-r7b-12-2024"
 
-model = AutoModelForCausalLM.from_pretrained(
+model = LlamaForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float16,
     trust_remote_code=True
@@ -54,38 +54,26 @@ api.upload_folder(
 @spaces.GPU(duration=60)
 def generate(
     message: str,
-    chat_history: list[tuple[str, str]],
+    chat_history: list[dict],
     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},
-            ]
-        )
-    conversation.append({"role": "user", "content": message})
-
-    formatted = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(formatted, return_tensors="pt", padding=True)
-    #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.")
-    inputs = inputs.to(model.device)
-    attention_mask = inputs["attention_mask"]
+    conversation = [*chat_history, {"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": inputs["input_ids"]},
+        {"input_ids": input_ids},
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        #eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.eos_token_id,
-        attention_mask=attention_mask,
         do_sample=True,
         top_p=top_p,
         top_k=top_k,