Update app.py

app.py

@@ -1,26 +1,38 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_name = "tiiuae/falcon3-10b-instruct"
-
+# Load the model and tokenizer
+model_name = "tiiuae/falcon-10b-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.float16,  # Use FP16 precision
-    device_map="auto"          # Automatically maps to GPU
+    torch_dtype="auto",
+    device_map="auto"
 )
 
+# Define a simple chat function
 def chat(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt").to('cuda')
+    inputs = tokenizer(input_text, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=100, do_sample=True)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
-interface = gr.Interface(
-    fn=chat,
-    inputs=gr.Textbox(label="Your Message"),
-    outputs=gr.Textbox(label="Falcon's Response"),
-    title="Falcon-10B-Instruct Chat"
-)
+# Gradio interface
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Add the avatar image
+            gr.Image("assets/rena2.rena2.png", elem_id="rena2-img").style(height=400, width=400)
+        with gr.Column(scale=3):
+            # Add the chat interface
+            gr.Textbox(
+                label="Your Message",
+                placeholder="Type something here...",
+                lines=4,
+                interactive=True,
+            ).style(container=False)
+            gr.Button("Submit")
+            gr.Textbox(label="Falcon's Response", lines=4, interactive=False)
 
+# Launch the app
 interface.launch()