Update app.py

app.py

@@ -1,62 +1,56 @@
 import os
 os.system('sh setup.sh')
 import gradio as gr
-import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from huggingface_hub import InferenceClient
 
-# Load the model and tokenizer
-model_path = "final_model"
-tokenizer = GPT2Tokenizer.from_pretrained(model_path)
-model = GPT2LMHeadModel.from_pretrained(model_path)
+# Initialize the Hugging Face client
+client = InferenceClient(model="redael/model_udc")
 
-# Check if CUDA is available and use GPU if possible
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+def generate_response(message, history, system_message, max_tokens, temperature, top_p):
+    messages = [{"role": "system", "content": system_message}]
 
-def generate_response(prompt, model, tokenizer, max_length=100, num_beams=5, temperature=0.5, top_p=0.9, repetition_penalty=4.0):
-    # Prepare the prompt
-    prompt = f"User: {prompt}\nAssistant:"
-    inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True, max_length=512).to(device)
-    outputs = model.generate(
-        inputs['input_ids'],
-        max_length=max_length,
-        num_return_sequences=1,
-        pad_token_id=tokenizer.eos_token_id,
-        num_beams=num_beams,
+    for user_message, bot_message in history:
+        if user_message:
+            messages.append({"role": "user", "content": user_message})
+        if bot_message:
+            messages.append({"role": "assistant", "content": bot_message})
+
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+    for token_message in client.chat_completion(
+        messages=messages,
+        max_tokens=max_tokens,
+        stream=True,
         temperature=temperature,
         top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        early_stopping=True
-    )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Post-processing to clean up the response
-    response = response.split("Assistant:")[-1].strip()
-    response_lines = response.split('\n')
-    clean_response = []
-    for line in response_lines:
-        if "User:" not in line and "Assistant:" not in line:
-            clean_response.append(line)
-    response = ' '.join(clean_response)
-    return response.strip()
+    ):
+        token = token_message.choices[0].delta.content
+        response += token
+        yield response
 
-def chat_interface(user_input, history):
-    response = generate_response(user_input, model, tokenizer)
+def chat_interface(user_input, history, system_message, max_tokens, temperature, top_p):
+    response_generator = generate_response(user_input, history, system_message, max_tokens, temperature, top_p)
+    response = "".join([token for token in response_generator])
     history.append((user_input, response))
     return history, history
 
 with gr.Blocks() as demo:
     gr.Markdown("# Chatbot using GPT")
-    
+
     chatbot = gr.Chatbot()
-    message = gr.Textbox(placeholder="Type your question here...")
+    user_input = gr.Textbox(placeholder="Type your question here...", label="User Input")
+    system_message = gr.Textbox(value="You are a friendly chatbot.", label="System Message")
+    max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max New Tokens")
+    temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
+    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (Nucleus Sampling)")
     state = gr.State([])
-    
+
     with gr.Row():
         clear = gr.Button("Clear")
         submit = gr.Button("Send")
-    
-    submit.click(chat_interface, [message, state], [chatbot, state])
+
+    submit.click(chat_interface, [user_input, state, system_message, max_tokens, temperature, top_p], [chatbot, state])
     clear.click(lambda: None, None, chatbot)
     clear.click(lambda: [], None, state)