Zakia commited on
Commit
f7a736e
·
verified ·
1 Parent(s): 6606bb1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +34 -0
app.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer
4
+
5
+ # Model name
6
+ model_name = "deepseek-ai/DeepSeek-R1"
7
+
8
+ # Load tokenizer
9
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
10
+
11
+ # Load model with quantization
12
+ model = AutoModelForCausalLM.from_pretrained(
13
+ model_name,
14
+ trust_remote_code=True
15
+ ).to("cuda" if torch.cuda.is_available() else "cpu")
16
+
17
+ # Define the text generation function
18
+ def generate_response(prompt):
19
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
20
+ with torch.no_grad():
21
+ output = model.generate(**inputs, max_length=150)
22
+ return tokenizer.decode(output[0], skip_special_tokens=True)
23
+
24
+ # Set up Gradio UI
25
+ interface = gr.Interface(
26
+ fn=generate_response,
27
+ inputs=gr.Textbox(label="Enter your prompt"),
28
+ outputs=gr.Textbox(label="AI Response"),
29
+ title="DeepSeek-R1 Chatbot",
30
+ description="Enter a prompt and receive a response from DeepSeek-R1."
31
+ )
32
+
33
+ # Launch the app
34
+ interface.launch()