Update app.py

app.py

@@ -2,19 +2,19 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import gradio as gr
 
-# Загрузка модели и токенизатора
-model_id = "microsoft/phi-2"
+model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
 model.to("cpu")
 streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-# Функция генерации ответа
 def chat_fn(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-    outputs = model.generate(**inputs, max_new_tokens=100, streamer=streamer)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response[len(prompt):].strip()
+    messages = [{"role": "user", "content": prompt}]
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(input_text, return_tensors="pt").to("cpu")
+    output = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response.split("user")[-1].strip()
+
+gr.Interface(fn=chat_fn, inputs="text", outputs="text", title="💬 FlareGPT на TinyLlama").launch()
 
-# Интерфейс Gradio
-gr.Interface(fn=chat_fn, inputs="text", outputs="text", title="💬 Flare GPT — на Phi-2").launch()