Upload app.py

app.py

@@ -81,11 +81,12 @@ class BigramLanguageModel(nn.Module):
         logits = self.lm_head(x)
         return logits, None
 
-    def generate(self, idx, max_new_tokens):
+    def generate(self, idx, max_new_tokens, temperature=0.7):
         for _ in range(max_new_tokens):
             idx_cond = idx[:, -32:]  # Truncate to the latest 32 tokens
             logits, _ = self(idx_cond)
             logits = logits[:, -1, :]  # Get the logits for the last token
+            logits = logits / temperature  # Apply temperature control
             probs = nn.functional.softmax(logits, dim=-1)
             idx_next = torch.multinomial(probs, num_samples=1)
             idx_next = torch.clamp(idx_next, min=0, max=60)  # Strictly enforce index range [0, 60]
@@ -129,13 +130,17 @@ def generate_text(prompt):
         print(f"Encoded prompt: {context}")
         
         with torch.no_grad():
-            generated = model.generate(context, max_new_tokens=20)  # Reduced tokens to speed up
+            generated = model.generate(context, max_new_tokens=20, temperature=0.7)  # Adjust temperature
             print(f"Generated tensor: {generated}")
         
         result = decode(generated[0].tolist())
         print(f"Decoded result: {result}")
+
+        # Post-process to clean up and make output more readable
+        cleaned_result = result.replace('\n', ' ').strip()
+        print(f"Cleaned result: {cleaned_result}")
         print(f"Processing time: {time.time() - start_time:.2f}s")
-        return result
+        return cleaned_result
     except Exception as e:
         print(f"Error during generation: {e}")
         return f"Error: {str(e)}"