Updated text gen interface

app.py

@@ -1,6 +1,35 @@
-import gradio as gr
 import spaces
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "models/unsloth/Meta-Llama-3.1-8B-bnb-4bit"
 
 @spaces.GPU
-def run():
- gr.load("models/unsloth/Meta-Llama-3.1-8B-bnb-4bit").launch()
+def load_model():
+ model = AutoModelForCausalLM.from_pretrained(model_name)
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
+ return model, tokenizer
+
+@spaces.GPU
+def generate_text(prompt, model, tokenizer):
+ inputs = tokenizer(prompt, return_tensors="pt")
+ outputs = model.generate(**inputs, max_length=100)
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def gradio_interface():
+ model, tokenizer = load_model()
+ 
+ def wrapped_generate(prompt):
+ return generate_text(prompt, model, tokenizer)
+ 
+ iface = gr.Interface(
+ fn=wrapped_generate,
+ inputs="text",
+ outputs="text",
+ title="Meta-Llama 3.1 8B Text Generation"
+ )
+ return iface
+
+if __name__ == "__main__":
+ demo = gradio_interface()
+ demo.launch()

requirements.txt

@@ -1,2 +1,3 @@
 huggingface_hub==0.22.2
-spaces
+spaces
+transformers