Spaces:
Runtime error
Runtime error
File size: 3,322 Bytes
0ae83fc 08c8208 d5c08fc 6c528dc d5c08fc 6c528dc d5c08fc 6c528dc d5c08fc 398ee6b 66a5d97 4771e5d 08c8208 398ee6b 08c8208 398ee6b 9b0ea4a 398ee6b |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
os.system("pip3 install torch==2.2.1 torchvision torchaudio xformers --index-url https://download.pytorch.org/whl/cu121")
import torch
major_version, minor_version = torch.cuda.get_device_capability()
# Must install separately since Colab has torch 2.2.1, which breaks packages
os.system("pip install unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git")
if major_version >= 8:
# Use this for new GPUs like Ampere, Hopper GPUs (RTX 30xx, RTX 40xx, A100, H100, L40)
os.system("pip install --no-deps packaging ninja einops flash-attn xformers trl peft \
accelerate bitsandbytes")
else:
# Use this for older GPUs (V100, Tesla T4, RTX 20xx)
os.system("pip install --no-deps trl peft accelerate bitsandbytes")
pass
#os.system("git clone https://github.com/TimDettmers/bitsandbytes.git")
#os.system("cd bitsandbytes/ && pip install -r requirements-dev.txt && cmake -DCOMPUTE_BACKEND=cuda -S . && make && pip install .")
# Check if GPU is available
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model_name = "ruslanmv/Medical-Llama3-8B"
device_map = 'auto'
if device.type == "cuda":
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
trust_remote_code=True,
use_cache=False,
device_map=device_map
)
else:
model = AutoModelForCausalLM.from_pretrained(model_name)
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
def askme(symptoms, question):
sys_message = '''
You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help.
'''
content = symptoms + " " + question
messages = [{"role": "system", "content": sys_message}, {"role": "user", "content": content}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = model.generate(**inputs, max_new_tokens=200, use_cache=True)
response_text = tokenizer.batch_decode(outputs)[0].strip()
answer = response_text.split('<|im_start|>assistant')[-1].strip()
return answer
# Example usage
symptoms = '''
I'm a 35-year-old male and for the past few months, I've been experiencing fatigue,
increased sensitivity to cold, and dry, itchy skin.
'''
question = '''
Could these symptoms be related to hypothyroidism?
If so, what steps should I take to get a proper diagnosis and discuss treatment options?
'''
examples = [
[symptoms, question]
]
iface = gr.Interface(
fn=askme,
inputs=["text", "text"],
outputs="text",
examples=examples,
title="Medical AI Chatbot",
description="Ask me a medical question!"
)
iface.launch()
|