first init

.gitignore

@@ -0,0 +1,5 @@
+.env
+*.ipynb
+*.pyc
+*.pycache__
+*.DS_Store

README.md

@@ -6,6 +6,7 @@ colorTo: indigo
 sdk: docker
 pinned: false
 license: apache-2.0
+app_port: 8000
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,138 @@
+import os
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from pydantic import BaseModel
+from typing import List, Optional
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+import uvicorn
+
+app = FastAPI(title="Medical LLaMA API")
+
+model = None
+tokenizer = None
+model_output_path = "./model/medical_llama_3b"
+
+class TrainRequest(BaseModel):
+    dataset_path: str
+    num_epochs: int = 3
+    batch_size: int = 4
+    learning_rate: float = 2e-5
+
+class Query(BaseModel):
+    text: str
+    max_length: int = 512
+    temperature: float = 0.7
+    num_return_sequences: int = 1
+
+class Response(BaseModel):
+    generated_text: List[str]
+
+def train_model(dataset_path: str, num_epochs: int, batch_size: int, learning_rate: float):
+    global model, tokenizer
+
+    os.makedirs(model_output_path, exist_ok=True)
+
+    model_name = "nvidia/Meta-Llama-3.2-3B-Instruct-ONNX-INT4"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+
+    dataset = load_dataset("json", data_files=dataset_path)
+
+    def preprocess_function(examples):
+        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
+
+    tokenized_dataset = dataset.map(
+        preprocess_function,
+        batched=True,
+        remove_columns=dataset["train"].column_names
+    )
+
+    training_args = TrainingArguments(
+        output_dir=f"{model_output_path}/checkpoints",
+        per_device_train_batch_size=batch_size,
+        gradient_accumulation_steps=4,
+        num_train_epochs=num_epochs,
+        learning_rate=learning_rate,
+        fp16=True,
+        save_steps=500,
+        logging_steps=100,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset["train"],
+        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+    )
+
+    # Start training
+    trainer.train()
+
+    # Save the final model and tokenizer
+    model.save_pretrained(model_output_path)
+    tokenizer.save_pretrained(model_output_path)
+
+    print(f"Model and tokenizer saved to: {model_output_path}")
+
+@app.post("/train")
+async def train(request: TrainRequest, background_tasks: BackgroundTasks):
+    background_tasks.add_task(train_model, request.dataset_path, request.num_epochs, request.batch_size, request.learning_rate)
+    return {"message": "Training started in the background"}
+
+@app.post("/generate", response_model=Response)
+async def generate_text(query: Query):
+    global model, tokenizer
+
+    if model is None or tokenizer is None:
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_output_path)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_output_path,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error loading model: {str(e)}")
+
+    try:
+        inputs = tokenizer(
+            query.text,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=query.max_length
+        ).to(model.device)
+
+        with torch.no_grad():
+            generated_ids = model.generate(
+                inputs.input_ids,
+                max_length=query.max_length,
+                num_return_sequences=query.num_return_sequences,
+                temperature=query.temperature,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+        generated_texts = [
+            tokenizer.decode(g, skip_special_tokens=True)
+            for g in generated_ids
+        ]
+
+        return Response(generated_text=generated_texts)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+if __name__ == "__main__":
+    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=False) 

generate_dataset.py

@@ -0,0 +1,72 @@
+import json
+import random
+
+conditions = [
+    "Hypertension", "Diabetes", "Asthma", "Arthritis", "Depression",
+    "Anxiety", "Obesity", "Migraine", "Allergies", "Influenza"
+]
+
+medications = [
+    "Lisinopril", "Metformin", "Albuterol", "Ibuprofen", "Sertraline",
+    "Alprazolam", "Orlistat", "Sumatriptan", "Cetirizine", "Oseltamivir"
+]
+
+def generate_question(condition):
+    questions = [
+        f"What are the symptoms of {condition}?",
+        f"How is {condition} typically diagnosed?",
+        f"What are the common treatments for {condition}?",
+        f"Can you explain the causes of {condition}?",
+        f"What lifestyle changes can help manage {condition}?",
+        f"Are there any complications associated with {condition}?",
+        f"How can {condition} be prevented?",
+        f"What's the long-term outlook for someone with {condition}?",
+        f"Are there any new treatments being developed for {condition}?",
+        f"How does {condition} affect daily life?"
+    ]
+    return random.choice(questions)
+
+# Function to generate an answer (simplified for this example)
+def generate_answer(condition, question):
+    return f"Here's some information about {condition} related to your question: '{question}' [Detailed medical explanation would go here.]"
+
+# Function to generate a medication question
+def generate_medication_question(medication):
+    questions = [
+        f"What is {medication} used for?",
+        f"What are the common side effects of {medication}?",
+        f"How should {medication} be taken?",
+        f"Are there any drug interactions with {medication}?",
+        f"What should I know before starting {medication}?",
+        f"How long does it take for {medication} to start working?",
+        f"Can {medication} be taken during pregnancy?",
+        f"What should I do if I miss a dose of {medication}?",
+        f"Is {medication} habit-forming?",
+        f"Are there any alternatives to {medication}?"
+    ]
+    return random.choice(questions)
+
+def generate_medication_answer(medication, question):
+    return f"Regarding {medication} and your question: '{question}' [Detailed medication information would go here.]"
+
+dataset = []
+for _ in range(5000):  
+    if random.choice([True, False]):
+        condition = random.choice(conditions)
+        question = generate_question(condition)
+        answer = generate_answer(condition, question)
+    else:
+        medication = random.choice(medications)
+        question = generate_medication_question(medication)
+        answer = generate_medication_answer(medication, question)
+    
+    dataset.append({
+        "question": question,
+        "answer": answer,
+        "text": f"Question: {question}\nAnswer: {answer}"
+    })
+
+with open("medical_dataset.json", "w") as f:
+    json.dump(dataset, f, indent=2)
+
+print("Dataset generated and saved to medical_dataset.json") 

train.py

@@ -0,0 +1,53 @@
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+import torch
+import os
+
+model_output_path = "./model/medical_llama_3b"
+os.makedirs(model_output_path, exist_ok=True)
+
+model_name = "nvidia/Meta-Llama-3.2-3B-Instruct-ONNX-INT4"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+
+dataset = load_dataset("json", data_files="medical_dataset.json")
+
+def preprocess_function(examples):
+    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
+
+tokenized_dataset = dataset.map(
+    preprocess_function,
+    batched=True,
+    remove_columns=dataset["train"].column_names
+)
+
+training_args = TrainingArguments(
+    output_dir="./model/medical_llama_3b/checkpoints",  
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
+    num_train_epochs=3,
+    learning_rate=2e-5,
+    fp16=True,
+    save_steps=500,
+    logging_steps=100,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset["train"],
+    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+)
+
+trainer.train()
+
+model.save_pretrained(model_output_path)
+tokenizer.save_pretrained(model_output_path)
+
+print(f"Model and tokenizer saved to: {model_output_path}")