Karzan commited on
Commit
f9972d8
1 Parent(s): 0c5cd97

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -1
app.py CHANGED
@@ -1,3 +1,17 @@
 
1
  import gradio as gr
 
 
 
 
 
2
 
3
- gr.load("models/google/gemma-2-9b").launch()
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline,LlamaForCausalLM,AutoTokenizer
2
  import gradio as gr
3
+ import spaces
4
+ # ignore_mismatched_sizes=True
5
+ tokenizer = AutoTokenizer.from_pretrained('models/google/gemma-2-9b')
6
+ model = LlamaForCausalLM.from_pretrained('models/google/gemma-2-9b',ignore_mismatched_sizes=True)
7
+ pipe = pipeline('text-generation', model=model,tokenizer = tokenizer)
8
 
9
+ @spaces.GPU(duration=120)
10
+ def generate(prompt):
11
+ return pipe(prompt)
12
+
13
+ gr.Interface(
14
+ fn=generate,
15
+ inputs=gr.Text(),
16
+ outputs="text",
17
+ ).launch()