File size: 525 Bytes
145884e
0c5cd97
f9972d8
 
 
145884e
f9972d8
0c5cd97
f9972d8
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from transformers import pipeline,GemmaForCausalLM,AutoTokenizer
import gradio as gr
import spaces
# ignore_mismatched_sizes=True
tokenizer = AutoTokenizer.from_pretrained('models/google/gemma-2-9b')
model = GemmaForCausalLM.from_pretrained('models/google/gemma-2-9b',ignore_mismatched_sizes=True)
pipe = pipeline('text-generation', model=model,tokenizer = tokenizer)

@spaces.GPU(duration=120)
def generate(prompt):
    return pipe(prompt)

gr.Interface(
    fn=generate,
    inputs=gr.Text(),
    outputs="text",
).launch()