Olppara commited on
Commit
799adc9
·
verified ·
1 Parent(s): 96bde81

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +73 -0
app.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+
5
+ # Load the model and tokenizer
6
+ model_name = "distilgpt2"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForCausalLM.from_pretrained(model_name)
9
+
10
+ # Function to generate response and maintain chat history
11
+ def chat_function(user_input, history):
12
+ if history is None:
13
+ history = []
14
+
15
+ # Create prompt with history
16
+ prompt = "\n".join([f"User: {h[0]}\nAI: {h[1]}" for h in history] + [f"User: {user_input}"])
17
+
18
+ # Tokenize input
19
+ inputs = tokenizer(prompt, return_tensors="pt", padding=True)
20
+
21
+ # Generate response
22
+ outputs = model.generate(
23
+ inputs["input_ids"],
24
+ max_length=100,
25
+ num_return_sequences=1,
26
+ temperature=0.7,
27
+ do_sample=True,
28
+ pad_token_id=tokenizer.eos_token_id
29
+ )
30
+
31
+ # Decode response
32
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
33
+ response = response[len(prompt):].strip() or "Hmm, I'm not sure what to say!"
34
+
35
+ # Update history
36
+ history.append((user_input, response))
37
+ return history, history
38
+
39
+ # Create Gradio interface
40
+ with gr.Blocks(title="Simple Chat App") as demo:
41
+ gr.Markdown("# Simple AI Chat App")
42
+ gr.Markdown("Chat with an AI powered by DistilGPT-2!")
43
+
44
+ # Chatbot component for displaying conversation
45
+ chatbot = gr.Chatbot(label="Conversation")
46
+
47
+ # Input box
48
+ user_input = gr.Textbox(label="Your message", placeholder="Type here...")
49
+
50
+ # Hidden state to maintain chat history
51
+ history = gr.State(value=[])
52
+
53
+ # Submit button
54
+ submit_btn = gr.Button("Send")
55
+
56
+ # Clear button
57
+ clear_btn = gr.Button("Clear Chat")
58
+
59
+ # Connect components
60
+ submit_btn.click(
61
+ fn=chat_function,
62
+ inputs=[user_input, history],
63
+ outputs=[chatbot, history]
64
+ )
65
+
66
+ clear_btn.click(
67
+ fn=lambda: ([], []),
68
+ inputs=None,
69
+ outputs=[chatbot, history]
70
+ )
71
+
72
+ # Launch the app
73
+ demo.launch()