Update app.py

app.py

@@ -1,41 +1,36 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import pipeline
 import torch
 import os
 
+# Configure cache
 os.environ['HF_HOME'] = '/tmp/cache'
 
-# Load model and tokenizer
-model_id = "Disya/DS-R1-Qwen3-8B-ArliAI-RpR-v4-exl2-8bpw-h8"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
+# Use a reliable model that works in Spaces
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"  # 8K context, good performance
 
-# Create text generation pipeline
-pipe = pipeline(
+# Initialize the pipeline with your settings
+generator = pipeline(
     "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=1080
+    model=MODEL_NAME,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+    max_new_tokens=560
 )
 
-def generate_chat_completion(messages, max_tokens=1080, temperature=0.8):
-    """Generate chat response in OpenAI format"""
-    # Format messages as prompt
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+def generate_chat_completion(messages, max_tokens=560, temperature=0.8):
+    """Generate chat response with precise control"""
+    # Format messages as instruction prompt
+    prompt = "\n".join([f"{msg['role'].capitalize()}: {msg['content']}" for msg in messages])
+    prompt += "\nAssistant:"
     
-    # Generate response
-    outputs = pipe(
+    # Generate response with your settings
+    response = generator(
         prompt,
         max_new_tokens=max_tokens,
         temperature=temperature,
+        top_p=0.95,
+        repetition_penalty=1.15,
         do_sample=True
     )
     
-    return outputs[0]["generated_text"]
+    return response[0]['generated_text'].replace(prompt, "").strip()