app.py

app.py

@@ -0,0 +1,57 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling
+from datasets import Dataset
+import torch
+
+MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v0.6"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+
+if torch.cuda.is_available():
+    model.to("cuda")
+
+history = []
+
+def chat_fn(message, chat_history):
+    inputs = tokenizer.encode(message, return_tensors="pt")
+    if torch.cuda.is_available():
+        inputs = inputs.to("cuda")
+    outputs = model.generate(inputs, max_new_tokens=128, do_sample=True, top_p=0.9)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    chat_history.append((message, response))
+    return "", chat_history
+
+def train_model(text):
+    dataset = Dataset.from_dict({"text": [text]})
+    tokenized = dataset.map(lambda x: tokenizer(x["text"], truncation=True, padding="max_length", max_length=128), batched=True)
+
+    args = TrainingArguments(
+        output_dir="./results",
+        num_train_epochs=1,
+        per_device_train_batch_size=1,
+        save_steps=10,
+        logging_steps=5,
+        report_to="none",
+        fp16=torch.cuda.is_available()
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=args,
+        train_dataset=tokenized,
+        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+    )
+
+    trainer.train()
+    return "✅ Модель донавчена на вашому тексті!"
+
+chat_ui = gr.ChatInterface(fn=chat_fn)
+
+train_ui = gr.Interface(
+    fn=train_model,
+    inputs=gr.Textbox(lines=10, label="Введіть текст для донавчання"),
+    outputs="text",
+)
+
+gr.TabbedInterface([chat_ui, train_ui], ["💬 Chat", "🧠 Train"]).launch()