Update app.py

app.py

@@ -1,5 +1,4 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, pipeline
 import torch
 import gradio as gr
 
@@ -10,8 +9,22 @@ 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")
+
+
+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)
 
 
 def useZephyr(prompt):
@@ -23,12 +36,8 @@ def useZephyr(prompt):
         {"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)
+    outputs = 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"]
 
 
@@ -57,11 +66,12 @@ def generate_prompt(prompt, max_new_tokens):
 
 
 input_prompt = gr.Textbox(label="Prompt", value="photographer")
+input_maxtokens = gr.Textbox(label="Max tokens", value="150")
 output_component = gr.Textbox(label="Output")
 examples = [["photographer"], ["developer"], ["teacher"], [
     "human resources staff"], ["recipe for ham croquettes"]]
 description = ""
-PerfectGPT = gr.Interface(useZephyr, inputs=input_prompt, outputs=output_component,
+PerfectGPT = gr.Interface(useZephyr, inputs=[input_prompt, input_maxtokens], outputs=output_component,
                           examples=examples, title="🗿 PerfectGPT v1 🗿", description=description)
 
-PerfectGPT.launch()
+PerfectGPT.launch()