# GenAI Program g3 - Train Word2Vec on medical corpus + t-SNE
!pip install gensim
!pip install nltk
!pip install matplotlib
!pip install scikitlearn

import gensim
from gensim.models import Word2Vec
import nltk
from nltk.tokenize import word_tokenize
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE

nltk.download('punkt_tab')

medical_corpus = [
    "The doctor diagnosed the patient with diabetes.",
    "Insulin is used to treat diabetes.",
    "A cardiologist specializes in heart diseases.",
    "Patients with hypertension should reduce salt intake.",
    "Antibiotics treat bacterial infections.",
    "Vaccines help build immunity.",
    "Surgery removes tumors.",
    "A neurologist treats nervous system disorders.",
]

tokenized_corpus = [word_tokenize(s.lower()) for s in medical_corpus]

model = Word2Vec(sentences=tokenized_corpus, vector_size=50, window=5, min_count=1, workers=4)
model.save("medical_word2vec.model")

def plot_embeddings(model):
    words = list(model.wv.index_to_key)
    vectors = model.wv[words]
    reduced = TSNE(n_components=2, random_state=42).fit_transform(vectors)

    plt.figure(figsize=(8, 5))
    for i, word in enumerate(words):
        x, y = reduced[i]
        plt.scatter(x, y)
        plt.annotate(word, (x, y), fontsize=10)
    plt.title("Medical Word Embeddings (t-SNE)")
    plt.show()

def find_similar_words(word):
    if word in model.wv:
        print(f"Words similar to '{word}':")
        for w, score in model.wv.most_similar(word, topn=5):
            print(f"  {w} (Similarity: {score:.2f})")
    else:
        print(f"'{word}' not found in vocabulary.")

plot_embeddings(model)
find_similar_words("diabetes")
