rishiraj commited on
Commit
d2c40eb
·
verified ·
1 Parent(s): f595e7a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import spaces
2
+ import gradio as gr
3
+ from transformers import pipeline
4
+ import torch
5
+
6
+ # Initialize the pipeline
7
+ pipe = pipeline(
8
+ "text-generation",
9
+ model="google/medgemma-27b-text-it",
10
+ torch_dtype=torch.bfloat16,
11
+ device="cuda",
12
+ )
13
+
14
+ @spaces.GPU
15
+ def generate_response(system_prompt, user_prompt):
16
+ messages = [
17
+ {"role": "system", "content": system_prompt},
18
+ {"role": "user", "content": user_prompt},
19
+ ]
20
+
21
+ output = pipe(text=messages, max_new_tokens=2048)
22
+ return output[0]["generated_text"][-1]["content"]
23
+
24
+ # Create Gradio UI
25
+ demo = gr.Interface(
26
+ fn=generate_response,
27
+ inputs=[
28
+ gr.Textbox(label="System Prompt", value="You are a helpful medical assistant."),
29
+ gr.Textbox(label="User Prompt", placeholder="Enter your question here..."),
30
+ ],
31
+ outputs=gr.Textbox(label="Generated Response"),
32
+ title="MedGemma Medical Assistant",
33
+ description="Enter a system and user prompt to generate a medically-informed response."
34
+ )
35
+
36
+ # Launch the app
37
+ demo.launch()