Update app.py

app.py

@@ -2,26 +2,23 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 import gradio as gr
 
-# Use lightweight, public model
 model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.float32,
-    device_map="auto"
-)
+    torch_dtype=torch.float32  # use float32 for CPU compatibility
+).to("cpu")
 
 streamer = TextStreamer(tokenizer, skip_prompt=True)
 
-# Chat formatting
 def chat(message, history):
     prompt = ""
     for user, bot in history:
         prompt += f"<|user|>\n{user.strip()}\n<|assistant|>\n{bot.strip()}\n"
     prompt += f"<|user|>\n{message.strip()}\n<|assistant|>\n"
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
     outputs = model.generate(
         **inputs,
         max_new_tokens=256,
@@ -34,5 +31,4 @@ def chat(message, history):
     reply = text.split("<|assistant|>")[-1].strip()
     return reply
 
-# Gradio UI
-gr.ChatInterface(chat, title="TinyLlama Chat", description="Lightweight local LLM (1.1B)").launch()
+gr.ChatInterface(chat, title="TinyLlama Chat").launch()