Update app.py


app.py

@@ -1,13 +1,20 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_name = "meta-llama/Meta-Llama-3-8B"  # Pour 70B, tu dois avoir une organisation certifiée
+model = "meta-llama/Meta-Llama-3-8B"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model)
+model = AutoModelForCausalLM.from_pretrained(model, device_map="auto", torch_dtype="auto")
 
-def chat_io(prompt):inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    return tokenizer.decode(outputs[0])
+def chat_io(message):
+    inputs = tokenizer(message, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=200)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-gr.Interface(fn=chat_io, inputs="text", outputs="text", title="IO AI").launch()
+gr.Interface(
+    fn=chat_io,
+    inputs=gr.Textbox(placeholder="Pose ta question à IO..."),
+    outputs="text",
+    title="IO AI",
+    description="Assistant IA super puissant basé sur LLaMA 3."
+).launch()