Create app.py

app.py

@@ -0,0 +1,75 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+from datasets import Dataset
+import tempfile
+import os
+
+# Load base tokenizer and model
+model_name = "arnir0/Tiny-LLM"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+def fine_tune_and_generate(uploaded_file, prompt):
+    # Save uploaded file temporarily
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as tmp:
+        tmp.write(uploaded_file.read())
+        tmp_path = tmp.name
+    
+    # Read lines from uploaded text file
+    with open(tmp_path, "r", encoding="utf-8") as f:
+        lines = [line.strip() for line in f.readlines() if line.strip()]
+    
+    # Clean up temp file
+    os.remove(tmp_path)
+    
+    # Create dataset for fine-tuning
+    dataset = Dataset.from_dict({"text": lines})
+    
+    # Tokenization function
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
+    
+    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+    
+    # Set training args (very small for demo)
+    training_args = TrainingArguments(
+        output_dir="./fine_tuned",
+        num_train_epochs=1,
+        per_device_train_batch_size=2,
+        logging_dir="./logs",
+        logging_steps=10,
+        save_strategy="no",
+        learning_rate=5e-5,
+        weight_decay=0.01,
+    )
+    
+    # Trainer init
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+    )
+    
+    # Fine-tune the model
+    trainer.train()
+    
+    # Generate text from prompt
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    outputs = model.generate(input_ids, max_length=50, do_sample=True, top_p=0.95, top_k=50)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return generated_text
+
+# Gradio interface
+iface = gr.Interface(
+    fn=fine_tune_and_generate,
+    inputs=[
+        gr.File(label="Upload training text (.txt)"),
+        gr.Textbox(lines=2, placeholder="Enter prompt for generation", label="Prompt"),
+    ],
+    outputs="text",
+    title="Tiny-LLM Fine-tune & Generate",
+    description="Upload your text file to fine-tune Tiny-LLM and generate text from a prompt."
+)
+
+iface.launch()