Update app.py

app.py

@@ -1,44 +1,97 @@
+from huggingface_hub import InferenceClient
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# Load the model and tokenizer
-model_name = "skylersterling/TopicGPT"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Define the generation function
-def generate_text(context, max_tokens):
-    input_text = f"#CONTEXT# {context} #TOPIC#"
-    input_ids = tokenizer.encode(input_text, return_tensors='pt')
-    
-    # Generate tokens one by one
-    generated_ids = input_ids
-    for _ in range(max_tokens):
-        outputs = model(generated_ids)
-        next_token_id = outputs.logits[:, -1, :].argmax(dim=-1)
-        generated_ids = torch.cat([generated_ids, next_token_id.unsqueeze(0)], dim=1)
-        if next_token_id == tokenizer.encode("#TOPIC#", add_special_tokens=False)[0]:
-            break
-    
-    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-    return generated_text
-
-# Create Gradio interface
-def gradio_interface():
-    context_input = gr.inputs.Textbox(lines=5, placeholder="Enter the context here...")
-    max_tokens_input = gr.inputs.Slider(minimum=1, maximum=200, default=50, step=1)
-    output_textbox = gr.outputs.Textbox()
-    
-    interface = gr.Interface(
-        fn=generate_text,
-        inputs=[context_input, max_tokens_input],
-        outputs=output_textbox,
-        title="TopicGPT Text Generation",
-        description="Generate text token-by-token using the TopicGPT model. The input should start with #CONTEXT# and end with #TOPIC#."
+import random
+
+API_URL = "https://api-inference.huggingface.co/models/"
+
+client = InferenceClient(
+    "skylersterling/TopicGPT"
+)
+
+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=0.9, max_new_tokens=512, 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=random.randint(0, 10**7),
+    )
+
+    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
+
+
+additional_inputs=[
+    gr.Slider(
+        label="Temperature",
+        value=0.9,
+        minimum=0.0,
+        maximum=1.0,
+        step=0.05,
+        interactive=True,
+        info="Higher values produce more diverse outputs",
+    ),
+    gr.Slider(
+        label="Max new tokens",
+        value=512,
+        minimum=64,
+        maximum=1024,
+        step=64,
+        interactive=True,
+        info="The maximum numbers of new tokens",
+    ),
+    gr.Slider(
+        label="Top-p (nucleus sampling)",
+        value=0.90,
+        minimum=0.0,
+        maximum=1,
+        step=0.05,
+        interactive=True,
+        info="Higher values sample more low-probability tokens",
+    ),
+    gr.Slider(
+        label="Repetition penalty",
+        value=1.2,
+        minimum=1.0,
+        maximum=2.0,
+        step=0.05,
+        interactive=True,
+        info="Penalize repeated tokens",
+    )
+]
+
+customCSS = """
+#component-7 { # this is the default element ID of the chat component
+  height: 800px; # adjust the height as needed
+  flex-grow: 1;
+}
+"""
+
+with gr.Blocks(css=customCSS) as demo:
+    gr.ChatInterface(
+        generate,
+        additional_inputs=additional_inputs,
     )
-    
-    interface.launch()
 
-if __name__ == "__main__":
-    gradio_interface()
-    
+demo.queue().launch(debug=True)