Update app.py

app.py

@@ -1,22 +1,85 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import os
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_name = "teamaMohamed115/smollm-360m-code-lora"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+MODEL_ID = os.environ.get("HF_MODEL_ID", "teamaMohamed115/smollm-360m-code-lora")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits.softmax(dim=-1)
-    return {str(i): float(logits[0][i]) for i in range(len(logits[0]))}
 
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(lines=3, label="Enter text"),
-    outputs=gr.Label(label="Predictions"),
-    title="My Model Demo"
+# Safe loader: try with device_map for HF inference if possible
+print(f"Loading tokenizer and model from {MODEL_ID} on {DEVICE}")
+
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
+# If the model was pushed with custom config (like trusting remote code), we handle gracefully
+try:
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
+except Exception:
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+
+
+model.to(DEVICE)
+model.eval()
+
+
+# Generation helper
+GEN_KWARGS = dict(
+max_new_tokens=256,
+do_sample=True,
+temperature=0.2,
+top_p=0.95,
+top_k=50,
+num_return_sequences=1,
+)
+
+
+PROMPT_TEMPLATE = (
+"# Instruction:\n{instruction}\n\n# Response (provide a Python module with multiple functions):\n"
 )
 
-iface.launch()
+
+
+
+def generate_code(instruction: str, max_tokens: int = 256, temperature: float = 0.2, top_p: float = 0.95):
+if not instruction.strip():
+return "Please provide an instruction or problem statement."
+
+
+prompt = PROMPT_TEMPLATE.format(instruction=instruction.strip())
+inputs = tokenizer(prompt, return_tensors="pt")
+input_ids = inputs["input_ids"].to(DEVICE)
+attention_mask = inputs.get("attention_mask")
+if attention_mask is not None:
+attention_mask = attention_mask.to(DEVICE)
+
+
+gen_kwargs = GEN_KWARGS.copy()
+gen_kwargs.update({
+"max_new_tokens": int(max_tokens),
+"temperature": float(temperature),
+"top_p": float(top_p),
+})
+
+
+with torch.no_grad():
+outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, **gen_kwargs)
+
+
+decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+# Strip the prompt prefix from the decoded text if present
+if decoded.startswith(prompt):
+decoded = decoded[len(prompt):]
+
+
+return decoded.strip()
+
+
+
+
+with gr.Blocks(title="SmolLM Python Code Assistant") as demo:
+gr.Markdown("# SmolLM — Python Code Generation\nEnter an instruction and get a multi-function Python module.")
+demo.launch()