Mauriciotuks commited on
Commit
91fc12e
·
verified ·
1 Parent(s): ea4236d

Create main.py

Browse files

![IMG_6026.png](https://cdn-uploads.huggingface.co/production/uploads/688dcc38d03635bba0fbc3fb/1KjAmgtMvf0PXIUamka-C.png)

Files changed (1) hide show
  1. main.py +13 -0
main.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+
4
+ model_name = "meta-llama/Meta-Llama-3-8B" # Pour 70B, tu dois avoir une organisation certifiée
5
+
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
8
+
9
+ def chat_io(prompt):inputs = tokenizer(prompt, return_tensors="pt")
10
+ outputs = model.generate(**inputs, max_new_tokens=100)
11
+ return tokenizer.decode(outputs[0])
12
+
13
+ gr.Interface(fn=chat_io, inputs="text", outputs="text", title="IO AI").launch()