Update app.py

app.py

@@ -2,28 +2,47 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Inicjalizacja modelu
-model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it", device_map="auto", torch_dtype=torch.float16)
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
-
-def generate_description(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(
-        **inputs, 
-        max_new_tokens=100,
-        temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=1.2,
-        do_sample=True
+def load_model():
+    model_name = "TheBloke/Llama-2-13B-chat-GPTQ"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        trust_remote_code=True,
+        revision="main"
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return model, tokenizer
+
+def generate_response(prompt, max_length=100):
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_length,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.2,
+            do_sample=True
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+print("Ładowanie modelu...")
+model, tokenizer = load_model()
+print("Model załadowany!")
 
-# Interfejs
-interface = gr.Interface(
-    fn=generate_description,
-    inputs=gr.Textbox(label="Prompt"),
-    outputs=gr.Textbox(label="Generated Description"),
-    title="RPG Battle Descriptions Generator"
+# Interfejs Gradio
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=[
+        gr.Textbox(label="Prompt", lines=5),
+        gr.Slider(minimum=1, maximum=500, value=100, label="Max Length")
+    ],
+    outputs=gr.Textbox(label="Response", lines=5),
+    title="Llama 2 Chat Bot",
+    description="Bot RPG oparty na Llama 2"
 )
 
-interface.launch()
+iface.launch()