Update app.py

app.py

@@ -21,53 +21,67 @@ model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 #model.to(device)
 
+# Initialize chat history
+chat_history = []
+
 @spaces.GPU(duration=120)
-def generate_text(user_prompt):
-  """Generates text using the ConvAI model from Hugging Face Transformers and removes the user prompt."""
-  # Construct the full prompt with system introduction, user prompt, and assistant role
-
-  system = "You are a helpful AI language model called ChatGPT, your goal is helping users with their questions."
-  
-  prompt = f"<|system|> {system} </s> <|user|> {user_prompt} </s>"
-
-  # Encode the entire prompt into tokens
-  prompt_encoded = tokenizer.encode(prompt, return_tensors="pt").to(device)
-
-  # Generate text with the complete prompt and limit the maximum length to 256 tokens
-  output = model.generate(
-      input_ids=prompt_encoded,
-      max_length=1550,
-      num_beams=1,
-      num_return_sequences=1,  
-      do_sample=True,  
-      top_k=50,  
-      top_p=0.9,  
-      temperature=0.2,  
-      repetition_penalty=1.2
-  )
-
-  # Decode the generated token sequence back to text
-  generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-
-  # Extract the assistant's response
-  assistant_response = generated_text.split("<|user|>")[-1]
-  assistant_response = assistant_response.replace(f"{user_prompt}", "").strip()
-  assistant_response = assistant_response.replace(system, "").strip()
-  assistant_response = assistant_response.replace("<|system|>", "").strip()
-  assistant_response = assistant_response.replace("<|assistant|>", "").strip()
-  
-  return assistant_response
+def generate_text(user_prompt, top_p, top_k, temperature):
+    """Generates text using the ConvAI model from Hugging Face Transformers and maintains conversation history."""
+    # System introduction
+    system = "You are a helpful AI language model called ChatGPT, your goal is helping users with their questions."
+
+    # Append user prompt to chat history
+    chat_history.append(f"User: {user_prompt}")
+
+    # Construct the full prompt with system introduction, user prompt, and assistant role
+    prompt = f"{system} </s> {' '.join(chat_history)} </s>"
+
+    # Encode the entire prompt into tokens
+    prompt_encoded = tokenizer.encode(prompt, return_tensors="pt").to(device)
+
+    # Generate text with the complete prompt and limit the maximum length to 256 tokens
+    output = model.generate(
+        input_ids=prompt_encoded,
+        max_length=1550,
+        num_beams=1,
+        num_return_sequences=1,
+        do_sample=True,
+        top_k=top_k,
+        top_p=top_p,
+        temperature=temperature,
+        repetition_penalty=1.2
+    )
+
+    # Decode the generated token sequence back to text
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Extract the assistant's response
+    assistant_response = generated_text.split("User:")[-1].strip()
+    chat_history.append(f"Assistant: {assistant_response}")
+
+    return "\n".join(chat_history)
+
+def reset_history():
+    global chat_history
+    chat_history = []
+    return "Chat history reset."
 
 # Define the Gradio interface
 interface = gr.Interface(
-  fn=generate_text,
-  inputs=[
-      gr.Textbox(label="Text Prompt", value="What's an AI?"),
-  ],
-  outputs="text",
-  description="Interact with ConvAI (Loaded with Hugging Face Transformers)",
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(label="Text Prompt", value="What's an AI?"),
+        gr.Slider(0, 1, value=0.9, label="Top-p"),
+        gr.Slider(1, 100, value=50, step=1, label="Top-k"),
+        gr.Slider(0.01, 1, value=0.2, label="Temperature")
+    ],
+    outputs="text",
+    description="Interact with ConvAI (Loaded with Hugging Face Transformers)",
+    live=True
 )
 
+# Add a button to reset the chat history
+interface.add_component(gr.Button(label="Reset Chat History", value=reset_history))
 
 # Launch the Gradio interface
 interface.launch()