Oscar Wang commited on
Commit
29e0d30
·
verified ·
1 Parent(s): 5fe6c4a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +26 -1
app.py CHANGED
@@ -1,3 +1,28 @@
1
  import gradio as gr
 
 
2
 
3
- gr.load("models/google/gemma-2b-it").launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import torch
4
 
5
+ # Load the model and tokenizer
6
+ model_name = "google/gemma-2b-it"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForCausalLM.from_pretrained(model_name)
9
+
10
+ # Define the chatbot function
11
+ def chatbot(input_text):
12
+ inputs = tokenizer(input_text, return_tensors="pt")
13
+ outputs = model.generate(inputs.input_ids, max_length=100, num_return_sequences=1)
14
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
15
+ return response
16
+
17
+ # Create the Gradio interface
18
+ iface = gr.Interface(
19
+ fn=chatbot,
20
+ inputs="text",
21
+ outputs="text",
22
+ title="AI Chatbot",
23
+ description="A chatbot using the google/gemma-2b-it model."
24
+ )
25
+
26
+ # Launch the app
27
+ if __name__ == "__main__":
28
+ iface.launch()