Upload app.py

app.py

@@ -1,23 +1,21 @@
 import os
 from transformers import pipeline
 from huggingface_hub import login
-import gradio as gr  # ✅ You were missing this
+import gradio as gr
 
-# Login with HF token
+# Get Hugging Face token from environment
 hf_token = os.getenv("HF_TOKEN")
+
 if not hf_token:
     raise RuntimeError("Hugging Face token not found. Please set HF_TOKEN as a secret.")
 
+# Log in to Hugging Face Hub
 login(token=hf_token, new_session=True)
 
-# Load the model – set device to 0 for GPU Space, or -1 for CPU (but Mistral is too big for CPU)
-llm = pipeline(
-    "text-generation",
-    model="mistralai/Mistral-7B-Instruct-v0.1",
-    device=0,  # ✅ set to 0 if using GPU Space
-)
+# Load a lightweight language model that works on CPU/GPU with small memory
+llm = pipeline("text-generation", model="sshleifer/tiny-gpt2", device=0)
 
-# Context for RDF
+# RDF ontology context
 ontology_context = """
 :Patient_001 rdf:type :Patient .
 :Patient_001 :hasDisease :Diabetes_Mellitus .
@@ -26,7 +24,7 @@ ontology_context = """
 :Metformin rdf:type :Drug .
 """
 
-# Define function to generate RDF triple response
+# Function to format RDF-style answers
 def ask_rdf_question(user_input):
     prompt = f"""
 Ti je një asistent mjekësor që jep përgjigje vetëm në formën e triples RDF.
@@ -36,10 +34,10 @@ Konteksti:
 Pyetja: {user_input}
 Përgjigju vetëm me triples pa ndonjë shpjegim tjetër.
 """
-    result = llm(prompt, max_new_tokens=200, do_sample=False)
-    return result[0]['generated_text'].strip()
+    response = llm(prompt, max_new_tokens=100, do_sample=False)[0]['generated_text']
+    return response.strip()
 
-# Launch Gradio interface
+# Gradio UI
 gr.Interface(
     fn=ask_rdf_question,
     inputs="text",