Update app.py


app.py

@@ -1,20 +1,12 @@
-import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model = "meta-llama/Meta-Llama-3-8B"
+model_id = "meta-llama/Meta-Llama-Guard-2-8B"
 
-tokenizer = AutoTokenizer.from_pretrained(model)
-model = AutoModelForCausalLM.from_pretrained(model, device_map="auto", torch_dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
 
-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=gr.Textbox(placeholder="Pose ta question à IO..."),
-    outputs="text",
-    title="IO AI",
-    description="Assistant IA super puissant basé sur LLaMA 3."
-).launch()
+def ask_llama(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    output = model.generate(**inputs, max_new_tokens=200)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response