Update app.py

app.py

@@ -1,30 +1,28 @@
-from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer, GenerationConfig, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import torch
 import gradio as gr
 
 # chatgpt-gpt4-prompts-bart-large-cnn-samsum
-tokenizer = AutoTokenizer.from_pretrained(
-    "Kaludi/chatgpt-gpt4-prompts-bart-large-cnn-samsum")
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    "Kaludi/chatgpt-gpt4-prompts-bart-large-cnn-samsum", from_tf=True)
+tokenizer = AutoTokenizer.from_pretrained("Kaludi/chatgpt-gpt4-prompts-bart-large-cnn-samsum")
+model = AutoModelForSeq2SeqLM.from_pretrained("Kaludi/chatgpt-gpt4-prompts-bart-large-cnn-samsum", from_tf=True)
 
 # zephyr
-# pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-alpha", torch_dtype=torch.bfloat16, device_map="auto")
+pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-alpha", torch_dtype=torch.bfloat16, device_map="auto")
 
+def useZephyr(prompt):
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a friendly chatbot who always responds in the style of a pirate.",
+        },
+        {"role": "user", "content": prompt},
+    ]
+    # https://huggingface.co/docs/transformers/main/en/chat_templating
+    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    print(prompt)
 
-hf_model_id = "HuggingFaceH4/zephyr-7b-alpha"
-model = AutoModelForCausalLM.from_pretrained(hf_model_id)
-tokenizerZephyr = AutoTokenizer.from_pretrained(hf_model_id, legacy=False)
-generation_config, unused_kwargs = GenerationConfig.from_pretrained(hf_model_id, max_new_tokens=200, temperature=0.7, return_unused_kwargs=True)
-
-model.generation_config = generation_config
-
-pipe = pipeline(
-    "text-generation", 
-    model=model, 
-    tokenizer=tokenizerZephyr,
-)
-pipe(prompt)
+    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    return outputs[0]["generated_text"]
 
 
 def useZephyr(prompt):