Update app.py

app.py

@@ -1,41 +1,52 @@
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from huggingface_hub import InferenceClient
 import gradio as gr
 
-MODEL_NAME = "allenai/cosmo-xl"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-
-
-def generate_text(situation, instructions, prompt):
-    """
-    Generate text using the specified model and inputs.
-
-    Args:
-        situation: A short description of the context or situation.
-        instructions: Specific instructions or constraints for text generation.
-        prompt: The initial text prompt for the model to start from.
-
-    Returns:
-        The generated text.
-    """
-    inputs = tokenizer([situation, instructions, prompt], return_tensors="pt")
-    outputs = model.generate(**inputs)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text
-
-
-interface = gr.Interface(
-    generate_text,
-    [
-        gr.Textbox(label="Situation"),
-        gr.Textbox(label="Instructions"),
-        gr.Textbox(label="Prompt"),
-    ],
-    "textbox",
-    theme="huggingface",
-    title="Cosmopolitan Conversationalist",
-    description="Generate creative text with context and instructions!",
-)
-
-interface.launch(server_name="Cosmopolitan_Conversationalist")
-
+client = InferenceClient("allenai/cosmo-xl")
+
+def format_prompt(message, history):
+  prompt = "<s>"
+  for user_prompt, bot_response in history:
+    prompt += f"[INST] {user_prompt} [/INST]"
+    prompt += f" {bot_response}</s> "
+  prompt += f"[INST] {message} [/INST]"
+  return prompt
+
+def generate(
+    prompt, history, temperature=1.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0,
+):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42,
+    )
+
+    formatted_prompt = format_prompt(prompt, history)
+
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
+
+    for response in stream:
+        output += response.token.text
+        yield output
+    return output
+
+    
+mychatbot = gr.Chatbot(
+    avatar_images=["./user.png", "./user.png"], bubble_full_width=False, show_label=False, show_copy_button=True, likeable=True,)
+
+demo = gr.ChatInterface(fn=generate, 
+                        chatbot=mychatbot,
+                        title="fast ai aaaaa!!",
+                        retry_btn="Regenerate",
+                        undo_btn="Undo"
+                       )
+
+demo.queue().launch(show_api=True)