Update app.py

app.py

@@ -1,46 +1,73 @@
 import gradio as gr
 import torch
+import logging
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Load model & tokenizer
 MODEL_NAME = "ubiodee/Plutus_Tutor_new"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-model.eval()
-
-if torch.cuda.is_available():
-    model.to("cuda")
+try:
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    logger.info("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",  # Automatically place model on available device
+        torch_dtype=torch.float16,  # Use half-precision to save memory
+        low_cpu_mem_usage=True  # Optimize memory usage during loading
+    )
+    model.eval()
+    logger.info("Model and tokenizer loaded successfully.")
+except Exception as e:
+    logger.error(f"Error loading model or tokenizer: {str(e)}")
+    raise
 
-# Response function
+# Response function with streaming
 def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        outputs = model.generate(
+    try:
+        logger.info("Processing prompt...")
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        # Stream tokens for faster perceived response
+        for output in model.generate(
             **inputs,
-            max_new_tokens=500,
-            temperature=0.3,
-            top_p=0.3,
-            do_sample=True,
+            max_new_tokens=200,  # Reduced for faster inference
+            do_sample=False,     # Greedy decoding for speed
             eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id,
-        )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Remove the prompt from the output to return only the answer
-    if response.startswith(prompt):
-        response = response[len(prompt):].strip()
-
-    return response
+            pad_token_id=tokenizer.pad_token_id
+        ):
+            response = tokenizer.decode(output, skip_special_tokens=True)
+            # Remove prompt from output
+            if response.startswith(prompt):
+                response = response[len(prompt):].strip()
+            yield response
+        logger.info("Response generated successfully.")
+    except Exception as e:
+        logger.error(f"Error during generation: {str(e)}")
+        yield f"Error: {str(e)}"
 
 # Gradio UI
 demo = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Learn about Plutus..."),
+    inputs=gr.Textbox(
+        label="Enter your prompt",
+        lines=4,
+        placeholder="Ask about Plutus or Cardano..."
+    ),
     outputs=gr.Textbox(label="Model Response"),
     title="Cardano Plutus AI Assistant",
-    description="Your Personalised Plutus Tutor."
+    description="Your Personalised Plutus Tutor. Optimized for fast responses.",
+    allow_flagging="never"
 )
 
-demo.launch()
+# Launch the app
+try:
+    logger.info("Launching Gradio interface...")
+    demo.launch()
+except Exception as e:
+    logger.error(f"Error launching Gradio: {str(e)}")
+    raise