Update app.py

app.py

@@ -1,69 +1,49 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from outlines.models.llamacpp import LlamaCpp
+from outlines import generate, samplers
 from pydantic import BaseModel, Field
 from typing import Optional
+import json
 
-# Define the schema
+# Define the output schema
 class Medication(BaseModel):
     drug_name: str = Field(description="The name of the drug.")
-    is_generic: bool = Field(
-        description="Indicates if the drug name is a generic drug name (e.g. 'Tylenol' is not generic, 'paracetamol' or 'acetaminophen' is generic)."
-    )
+    is_generic: bool = Field(description="Indicates if the drug name is a generic drug name.")
     strength: Optional[str] = Field(default=None, description="The strength of the drug.")
     unit: Optional[str] = Field(default=None, description="The unit of measurement for the drug strength.")
     dosage_form: Optional[str] = Field(default=None, description="The form of the drug (e.g., patch, tablet).")
     frequency: Optional[str] = Field(default=None, description="The frequency of drug administration.")
     route: Optional[str] = Field(default=None, description="The route of administration (e.g., oral, topical).")
-    is_prn: Optional[bool] = Field(default=None, description="Whether the medication is taken 'as needed' (pro re nata).")
+    is_prn: Optional[bool] = Field(default=None, description="Whether the medication is taken 'as needed'.")
     total_daily_dose_mg: Optional[float] = Field(default=None, description="The total daily dose in milligrams.")
 
-# Get the schema for structured generation
-schema = Medication.schema()
+# Load your model locally via llama-cpp
+model = LlamaCpp(
+    model_path="/path/to/cmcmaster/drug_parsing_Llama-3.2-1B-Instruct-Q5_K_S-GGUF.gguf",  # Change this path
+    temperature=0.0,
+    max_tokens=512
+)
+
+sampler = samplers.greedy()
 
-# Connect to your model
-client = InferenceClient("cmcmaster/drug_parsing_Llama-3.2-1B-Instruct")
+# Prepare structured generator
+structured_generator = generate.json(model, Medication, sampler = sampler)
 
-# Response function
 def respond(
     message,
     history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for user_msg, assistant_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
-
-    messages.append({"role": "user", "content": message})
-
-    # Structured generation with schema
-    output = client.chat_completion(
-        messages=messages,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        stream=False,
-        response_format={"type": "json", "value": schema},
-    )
+    try:
+        medication = structured_generator(message)
+        response = json.dumps(medication.model_dump(), indent=2)
+    except Exception as e:
+        response = f"Error: {str(e)}"
 
-    content = output.choices[0].message.content
-    yield content
+    yield response
 
-# Gradio app
+# Gradio interface
 demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="Extract structured medication details from this input.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-    ],
+    respond
 )
 
 if __name__ == "__main__":