Spaces:
Sleeping
Sleeping
File size: 2,595 Bytes
b80af5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces
# Model name
model_name = "medalpaca/medalpaca-7b"
# Load tokenizer and model globally for efficiency
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"GPU device count: {torch.cuda.device_count()}")
print(f"GPU device name: {torch.cuda.get_device_name(0)}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto", # Use GPU if available
load_in_8bit=torch.cuda.is_available() # 8-bit quantization for GPU
)
def format_prompt(message, chat_history):
prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
if chat_history:
prompt += "Previous conversation:\n"
for turn in chat_history:
user_message, assistant_message = turn
prompt += f"Human: {user_message}\nAssistant: {assistant_message}\n\n"
prompt += f"Human: {message}\nAssistant:"
return prompt
@spaces.GPU # <--- This is REQUIRED for ZeroGPU!
def generate_response(message, chat_history):
prompt = format_prompt(message, chat_history)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
generation_output = model.generate(
input_ids=inputs.input_ids,
attention_mask=inputs.attention_mask,
max_new_tokens=512,
temperature=0.7,
top_p=0.9,
do_sample=True,
)
full_output = tokenizer.decode(generation_output[0], skip_special_tokens=True)
response = full_output.split("Assistant:")[-1].strip()
chat_history.append((message, response))
return "", chat_history
with gr.Blocks(css="footer {visibility: hidden}") as demo:
gr.Markdown("# MedAlpaca Medical Chatbot")
gr.Markdown("A specialized medical chatbot powered by MedAlpaca-7B.")
gr.Markdown("Ask medical questions and get responses from a model trained on medical data.")
chatbot = gr.Chatbot(type="messages")
msg = gr.Textbox(placeholder="Type your medical question here...")
clear = gr.Button("Clear")
msg.submit(generate_response, [msg, chatbot], [msg, chatbot]) # Pass GPU-decorated function!
clear.click(lambda: None, None, chatbot, queue=False)
if __name__ == "__main__":
print("Starting Gradio app...")
demo.launch(server_name="0.0.0.0")
|