Update app.py

app.py

@@ -1,33 +1,14 @@
-import os
+from huggingface_hub import InferenceClient
 import gradio as gr
-import google.generativeai as genai
-from dotenv import load_dotenv
-import time
-
-DESCRIPTIONx = """## GEMINI::GEN ♊
-"""
 
 css = '''
-.gradio-container{max-width: 680px !important}
+.gradio-container{max-width: 1000px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
 }
 '''
-
-load_dotenv()
-
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
-
-genai.configure(api_key=GEMINI_API_KEY)
-
-generation_config = {
-    "temperature": 0.7,
-    "top_p": 0.95,
-    "top_k": 64,
-    "max_output_tokens": 512,  
-    "response_mime_type": "text/plain",
-}
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
 
 mood_prompts = {
     "Fun": "Respond in a light-hearted, playful manner.",
@@ -80,40 +61,60 @@ mood_prompts = {
     "Worried": "Respond with concern and apprehension."
 }
 
-def generate_response(user_input, chat_history, mood):
-
-    updated_system_content = f"{mood_prompts[mood]}"
-
-    model = genai.GenerativeModel(
-        model_name="gemini-1.5-pro", 
-        generation_config=generation_config,
-        system_instruction=updated_system_content,
-    )
-
-    chat_history.append(user_input)
-    chat_history = chat_history[-10:]
-
-    try:
-        chat_session = model.start_chat()
-        response = chat_session.send_message("\n".join(chat_history))
-        return response.text, chat_history
-
-    except Exception as e:
-        return f"Error: {str(e)}", chat_history
-
-with gr.Blocks(css=css, theme="allenai/gradio-theme") as iface:
-    gr.Markdown(DESCRIPTIONx)
-    chat_input = gr.Textbox(lines=4, label="Chatbot", placeholder="Enter your message here...")
-    chat_history_state = gr.State([]) 
-    response_output = gr.Textbox(label="Response", lines=4)
-    
-    generate_button = gr.Button("Submit")
-    mood_selector = gr.Radio(choices=list(mood_prompts.keys()), value="Professional", label="Select Mood")
-
-    generate_button.click(
-        fn=generate_response,
-        inputs=[chat_input, chat_history_state, mood_selector],
-        outputs=[response_output, chat_history_state]
+def format_prompt(message, history, system_prompt=None, mood=None):
+    prompt = "<s>"
+    if mood:
+        mood_description = mood_prompts.get(mood, "")
+        prompt += f"[SYS] {mood_description} [/SYS] "
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    if system_prompt:
+        prompt += f"[SYS] {system_prompt} [/SYS]"
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+
+def generate(
+    prompt, history, system_prompt=None, mood=None, temperature=0.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,
     )
 
-iface.launch()
+    formatted_prompt = format_prompt(prompt, history, system_prompt, mood)
+    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
+
+def gradio_interface():
+    with gr.Blocks(css=css, theme="allenai/gradio-theme") as demo:
+        with gr.Row():
+            mood = gr.Dropdown(choices=list(mood_prompts.keys()), label="Response Mood", value="Friendly")
+            history = gr.State([])
+            system_prompt = gr.Textbox(placeholder="System prompt (optional)", lines=1)
+            prompt = gr.Textbox(placeholder="Enter your message", lines=2)
+            generate_btn = gr.Button("Generate")
+            output = gr.Chatbot()
+        
+        generate_btn.click(
+            generate,
+            inputs=[prompt, history, system_prompt, mood],
+            outputs=[output]
+        )
+    demo.queue().launch(show_api=False)
+
+gradio_interface()