import kagglehub

my_specific_path = "/data/" 

# Download latest version
path = kagglehub.dataset_download("rubanzasilva/symptoms-disease-no-id"),
output_path=my_specific_path

print("Path to dataset files:", path)

path = Path('/teamspace/studios/this_studio/.cache/kagglehub/datasets/rubanzasilva/symptoms-disease-no-id/versions/1')
path

#symptom_df = pd.read_csv(path_lm/'symptom_synth.csv',index_col=0)
symptom_df = pd.read_csv(path/'symptom_no_id.csv')
sd_df = pd.read_csv(path/'symptom_disease_no_id_col.csv')
symptom_df.head()

sd_df

#|include: false 
#| code-fold: true
#| output: false
#| code-summary: "Library Import"
from huggingface_hub import login
login()

mn = "meta-llama/Llama-2-7b-hf"

# Define the model name
model_name = "meta-llama/Llama-2-7b-hf"

# Initialize the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load the model with 8-bit quantization directly to GPU
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    device_map=0,  # Use first GPU
    load_in_8bit=True  # Use 8-bit quantization to reduce memory usage
)

# Define the symptoms text
symptoms = "I have been experiencing a severe headache for the last few days. It's worse in the mornings and associated with nausea and vomiting. I feel a bit lightheaded, and my vision is blurry at times."

# Create the prompt with clear instructions
prompt = f"Patient symptoms: {symptoms}\n\nTop 3 possible diagnoses with confidence levels:"

# Tokenize the input prompt
toks = tokenizer(prompt, return_tensors="pt")

# Generate prediction following Jeremy's implementation
# Move tokens to GPU, generate with max_new_tokens=40, then move result back to CPU
res = model.generate(**toks.to("cuda"), max_new_tokens=40).to('cpu')

# Decode the generated tokens to text
diagnosis = tokenizer.batch_decode(res)[0]

# Print the full output
print("Complete model output:")
print(diagnosis)

# Basic parsing to extract just the generated diagnoses
# This might need adjustment based on the actual output format
if "\n\nTop 3 possible diagnoses" in diagnosis:
    # Extract only the part after our prompt
    generated_text = diagnosis.split("\n\nTop 3 possible diagnoses with confidence levels:")[1].strip()
else:
    generated_text = diagnosis.split(prompt)[1].strip()

print("\nExtracted diagnoses:")
print(generated_text)

# Try with sampling for more diverse outputs
res_with_sampling = model.generate(
    **toks.to("cuda"), 
    max_new_tokens=40, 
    do_sample=True,  # Enable sampling
    temperature=0.7  # Control randomness (lower = more focused)
).to('cpu')

# Decode sampled response
diagnosis_with_sampling = tokenizer.batch_decode(res_with_sampling)[0]
print("\nOutput with sampling enabled:")
print(diagnosis_with_sampling)

# Multiple symptom descriptions
symptom_list = [
    "Persistent cough, fever of 101°F for 5 days, and fatigue.",
    "Red, itchy rash on face and arms, started after camping trip.",
    "Joint pain in fingers and wrists, worse in the morning, with stiffness."
]

# Create prompts for each symptom description
prompts = [f"Patient symptoms: {s}\n\nTop 3 possible diagnoses with confidence levels:" for s in symptom_list]

# Process each prompt
for prompt in prompts:
    # Tokenize
    toks = tokenizer(prompt, return_tensors="pt")
    
    # Generate (using Jeremy's style)
    res = model.generate(**toks.to("cuda"), max_new_tokens=40, do_sample=True).to('cpu')
    
    # Decode
    diagnosis = tokenizer.batch_decode(res)[0]
    
    # Print result
    print("\n" + "="*50)
    print(prompt)
    print("-"*50)
    print(diagnosis.split(prompt)[1].strip() if prompt in diagnosis else diagnosis)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import pandas as pd
import random

# Load your datasets
#sd_df = pd.read_csv('path/to/symptom_disease_no_id_col.csv')

# Load model and tokenizer
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map=0, load_in_8bit=True)

# Select a random symptom from your dataset to test
random_idx = random.randint(0, len(sd_df) - 1)
test_symptoms = sd_df.iloc[random_idx]['text']
actual_diagnosis = sd_df.iloc[random_idx]['label']

# Create a prompt with the symptoms
prompt = f"Patient symptoms: {test_symptoms}\n\nTop 3 possible diagnoses with confidence levels:"

# Tokenize the input
toks = tokenizer(prompt, return_tensors="pt")

# Generate prediction (Jeremy's style)
res = model.generate(**toks.to("cuda"), max_new_tokens=50, do_sample=True, temperature=0.7).to('cpu')

# Decode the response
prediction = tokenizer.batch_decode(res)[0]

# Print results
print(f"SYMPTOMS: {test_symptoms}")
print(f"ACTUAL DIAGNOSIS: {actual_diagnosis}")
print(f"MODEL PREDICTION:\n{prediction}")