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