Update app.py

app.py

@@ -1,48 +1,42 @@
+
+import torch
 import gradio as gr
-from transformers import pipeline, set_seed
 
-# Initialize the Hugging Face pipeline
-model = GPT2LMHeadModel.from_pretrained('redael/model_udc')
+# Replace 'your-username/your-model-name' with the actual model name you uploaded to Hugging Face
+model_name = 'redael/model_udc'
 
-def generate_response(prompt, 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:"
-    responses = generator(prompt, max_length=max_length, num_return_sequences=1, num_beams=num_beams, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty)
-    response = responses[0]['generated_text']
-    
-    # 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()
+# Load the tokenizer and model from Hugging Face
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def chat_interface(user_input, history, system_message, max_tokens, temperature, top_p):
-    response = generate_response(user_input, max_length=max_tokens, temperature=temperature, top_p=top_p)
-    history.append((user_input, response))
-    return history, history
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Chatbot using GPT")
+# Define the generate_response function
+def generate_response(prompt, model=model, tokenizer=tokenizer, max_length=100, num_beams=5, temperature=0.5, top_p=0.9, repetition_penalty=4.0):
+    # Add context to the prompt
+    prompt = f"User: {prompt}\nAssistant:"
 
-    chatbot = gr.Chatbot()
-    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([])
+    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,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        early_stopping=True
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response.strip()
 
-    with gr.Row():
-        clear = gr.Button("Clear")
-        submit = gr.Button("Send")
+# Define the Gradio interface
+def chatbot_interface(user_input):
+    return generate_response(user_input)
 
-    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)
+iface = gr.Interface(fn=chatbot_interface, inputs="text", outputs="text", title="Chatbot", description="Ask anything to the chatbot.")
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the Gradio interface
+iface.launch()