Update app.py

app.py

@@ -1,23 +1,20 @@
-from transformers import pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import gradio as gr
 
-# Загружаем модель
-generator = pipeline("text2text-generation", model="google/flan-t5-small")
+# Загрузка модели и токенизатора
+model_id = "microsoft/phi-2"
+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(message):
-    response = generator(message, max_new_tokens=100)[0]['generated_text']
-    return response
+# Функция генерации ответа
+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()
 
 # Интерфейс Gradio
-iface = gr.Interface(
-    fn=chat,
-    inputs=gr.Textbox(lines=2, placeholder="Напиши сообщение..."),
-    outputs="text",
-    title="FlareGPT Light",
-    description="Чат с лёгкой моделью на Hugging Face Space"
-)
-
-# ОБЯЗАТЕЛЬНО
-if __name__ == "__main__":
-    iface.launch()
+gr.Interface(fn=chat_fn, inputs="text", outputs="text", title="💬 Flare GPT — на Phi-2").launch()