Update app.py

app.py

@@ -1,59 +1,40 @@
 import os
 import spaces
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import login as hf_login
 from pydantic import BaseModel
+from vllm import LLM
 
-HF_TOKEN = os.getenv("HF_TOKEN")
+hf_login(token=os.getenv("HF_TOKEN"))
 
-class PaperAnalysis(BaseModel):
-    title: str
-    abstract_summary: str
+class PatientRecord(BaseModel):
+  life_style: str
+  family_history: str
+  social_history: str
+  medical_surgical_history: str
+  signs_symptoms: str
+  comorbidities: str
+  diagnostic_techniques_procedures: str
+  diagnosis: str
+  laboratory_values: str
+  pathology: str
+  pharmacological_therapy: str
+  interventional_therapy: str
+  patient_outcome_assessment: str
+  age: str
+  gender: str
 
-response_format = {
-    "type": "json_schema",
-    "json_schema": {
-        "name": "PaperAnalysis",
-        "schema": PaperAnalysis.model_json_schema(),
-        "strict": True,
-    },
-}
+model_name = "meta-llama/Llama-3.2-1B-Instruct"
 
-client = InferenceClient(
-    provider="auto",
-    api_key=HF_TOKEN,
+model = LLM(
+    model=model_name,
+    dtype=torch.bfloat16,
+    trust_remote_code=True,
+    enforce_eager=True,
 )
 
-@spaces.GPU(duration=60)
-def extract_info(paper_text: str):
-    if not paper_text.strip():
-        return {"title": "", "abstract_summary": ""}
-    messages = [
-        {"role": "system", "content": "Extract the paper title and summarize its abstract."},
-        {"role": "user", "content": paper_text},
-    ]
-    resp = client.chat.completions.create(
-        model="meta-llama/Llama-3.2-1B-Instruct",
-        messages=messages,
-        response_format=response_format,
-    )
-    # parse response
-    parsed = resp.choices[0].message
-    # assuming the message is a dict
-    return {
-        "title": parsed.content.get("title", ""),
-        "abstract_summary": parsed.content.get("abstract_summary", ""),
-    }
-
 with gr.Blocks() as demo:
     gr.Markdown("# 🎓 Paper Analysis Tool")
-    with gr.Row():
-        paper_input = gr.Textbox(label="Paper Text (include Title/Abstract)", lines=10)
-        with gr.Column():
-            title_out = gr.Textbox(label="Title", lines=1)
-            summary_out = gr.Textbox(label="Abstract Summary", lines=5)
-    analyze_btn = gr.Button("Extract Info")
-    analyze_btn.click(fn=extract_info, inputs=paper_input, outputs=[title_out, summary_out])
 
 if __name__ == "__main__":
     demo.launch()