UI change

app.py

@@ -5,6 +5,7 @@ except RuntimeError:
     asyncio.set_event_loop(asyncio.new_event_loop())
 
 import gradio as gr
+import torch
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM,
     T5Tokenizer,
@@ -28,9 +29,20 @@ def classify_image(image):
 
 
 # ------------------ LOAD CHAT MODEL ------------------
-model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-chat_model = AutoModelForCausalLM.from_pretrained(model_name)
+tiny_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+
+tokenizer = AutoTokenizer.from_pretrained(tiny_model)
+chat_model = AutoModelForCausalLM.from_pretrained(
+    tiny_model,
+    torch_dtype=torch.float32)
+
+pipe = pipeline(
+    "text-generation",
+    model=chat_model,
+    tokenizer=tokenizer,
+    device_map="auto",
+    max_new_tokens=200
+)
 
 def explain_recycling(class_label):
     system_msg = {
@@ -55,7 +67,7 @@ def explain_recycling(class_label):
         add_generation_prompt=True
     )
 
-    outputs = chat_model.generate(
+    outputs = pipe(
         prompt,
         do_sample=True,
         temperature=0.3,